문제 23

위치 찾기

내 상태: 미제출
난이도: 23번 문제 난이도 보기

Silver IV
출제자: rlatjwls7882
시간 제한: 1000 ms
메모리 제한: 1024 MB

태그 보기

이분 탐색

정수 배열 $A_0, A_1, \cdots, A_{N-1}$이 있다.

$Q$개의 쿼리를 순서대로 처리하라.

각 쿼리마다 정수 $x$가 배열 $A$에 있다면 그 인덱스를 출력하고, 없다면 -1을 출력하라.

배열의 인덱스는 $0$부터 시작한다.

입력

첫 번째 줄에 두 정수 $N$, $Q$가 공백으로 구분되어 주어진다. $(1 \le N, Q \le 100\,000)$

두 번째 줄에 $N$개의 정수 $A_0, A_1, \cdots, A_{N-1}$이 공백으로 구분되어 주어진다. $(-10^9 \le A_i \le 10^9; A_0 \le A_1 \le \cdots \le A_{N-1})$

세 번째 줄부터 $Q$개의 줄에 걸쳐 정수 $x$가 하나씩 주어진다. $(-10^9 \le x \le 10^9)$

출력

각 쿼리마다 결과를 한 줄에 하나씩 출력한다.

예제 입력 1

5 6
-4 1 3 8 12
3
7
-4
12
0
1

예제 출력 1

공식 해설 풀이 스포일러 보기 접기

정렬된 배열에서는 $x$보다 작은 원소들이 앞에, $x$ 이상인 원소들이 뒤에 모여 있다. 이 두 부분의 경계를 이분 탐색한다.

경계 위치는 $0$부터 $N$까지이며 $N$은 $x$ 이상인 원소가 없다는 뜻이다. 처음 탐색 범위를 $[0,N]$으로 두고, 두 끝점이 다른 동안 중간 위치 $m$을 내림하여 잡는다. $A_m<x$이면 경계는 $m$보다 뒤에 있고, 아니라면 $m$ 또는 그 앞에 있다. 한 위치가 남을 때까지 범위를 줄이면 $x$ 이상인 첫 위치를 얻는다.

그 위치가 $N$이 아니면서 값이 $x$이면 해당 인덱스를 출력한다. 위치가 $N$이거나 값이 $x$보다 크면 배열에 $x$가 없으므로 $-1$을 출력한다. 중복값이 있다면 첫 등장 위치를 구하게 된다.

시간복잡도는 $O(N+Q\log(N+1))$, 공간복잡도는 $O(N)$이다.

구현 · 언어별 풀이

C++ 구현

vector<int>에 수열을 저장하고 algorithm의 lower_bound로 경계를 찾는다. 끝 반복자를 역참조하지 않도록 먼저 검사한다.

auto it = std::lower_bound(a.begin(), a.end(), x);
if (it != a.end() && *it == x)
    std::cout << it - a.begin() << '\n';
else
    std::cout << -1 << '\n';

C 구현

lo = 0, hi = n으로 시작한다. a[mid] < x이면 lo = mid + 1, 아니면 hi = mid로 갱신한다. 종료 후 lo < n && a[lo] == x이면 lo를, 아니라면 $-1$을 출력한다.

int lo = 0, hi = n;
while (lo < hi) {
    int mid = lo + (hi - lo) / 2;
    if (a[mid] < x) lo = mid + 1;
    else hi = mid;
}

Python 구현

bisect_left는 $x$ 이상인 첫 위치를 반환한다. 끝 인덱스인지 먼저 검사하고, 해당 값이 $x$인지 확인한다.

import sys
from bisect import bisect_left

read = sys.stdin.buffer.readline
n, q = map(int, read().split())
a = list(map(int, read().split()))
answers = []
for _ in range(q):
    x = int(read())
    i = bisect_left(a, x)
    answers.append(str(i if i < n and a[i] == x else -1))
sys.stdout.write("\n".join(answers))

Java 구현

a[mid] < x를 기준으로 lower bound를 직접 구하고 lo < n && a[lo] == x인지 확인한다.

int lo = 0, hi = n;
while (lo < hi) {
    int mid = lo + (hi - lo) / 2;
    if (a[mid] < x) lo = mid + 1;
    else hi = mid;
}

Rust 구현

slice.partition_point(|v| *v < x)로 첫 번째 $x$ 이상 위치를 얻는다. 결과가 배열 길이인지 확인한 뒤 해당 원소가 $x$인지 검사한다. 이 방법은 중복 값의 첫 위치를 반환한다.

let lo = a.partition_point(|&v| v < x);

JavaScript 구현

JavaScript에는 lower_bound가 없으므로 배열의 경계를 찾는 짧은 이분 탐색을 쓴다. 숫자 배열의 인덱스는 Number로 충분하며, 중간 위치는 Math.floor((l + r) / 2)로 구한다.

let l = 0, r = a.length;
while (l < r) {
    const mid = Math.floor((l + r) / 2);
    if (a[mid] < x) l = mid + 1;
    else r = mid;
}
const answer = l < a.length && a[l] === x ? l : -1;

끝 인덱스인지 먼저 확인하고 해당 값이 $x$이면 인덱스를 출력한다.

제출

언어