문제 56

k번째 공 찾기

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

Platinum V
출제자: rlatjwls7882
시간 제한: 1000 ms
메모리 제한: 512 MB

태그 보기

세그먼트 트리이분 탐색

$1$번부터 $N$번까지 번호가 붙은 상자가 있다.

$i$번 상자에는 처음에 $A_i$개의 공이 들어 있다.

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

쿼리는 다음 중 하나이다.

1 i x: $i$번 상자에 공 $x$개를 추가한다. $(1 \le i \le N; 1 \le x \le 10^9)$
2 k: $1$번 상자부터 $N$번 상자까지 순서대로 공을 나열했을 때 $k$번째 공이 들어 있는 상자의 번호를 출력한다. $(1 \le k \le S)$

여기서 $S$는 쿼리를 처리하기 전 모든 상자에 들어 있는 공의 개수의 합이다.

입력

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

두 번째 줄에 $N$개의 정수 $A_1, A_2, \cdots, A_N$이 공백으로 구분되어 주어진다. $(0 \le A_i \le 10^9)$

세 번째 줄부터 $Q$개의 줄에 걸쳐 쿼리가 하나씩 주어진다.

출력

2 k 쿼리마다 $k$번째 공이 들어 있는 상자의 번호를 한 줄에 하나씩 출력한다.

예제 입력 1

예제 출력 1

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

왼쪽 상자의 공들이 항상 먼저 나열된다. 따라서 한 구간의 공 수를 알면 찾는 공이 그 구간 안에 있는지 판단할 수 있다.

상자별 공 수를 리프에 놓고, 각 노드에 두 자식의 공 수 합을 저장하는 세그먼트 트리를 만든다. 공을 추가하면 해당 리프와 조상들의 합을 갱신한다.

$k$번째 공을 찾을 때는 루트에서 시작한다. 왼쪽 자식의 공 수를 $c$라 하면 $k\le c$일 때 왼쪽으로 내려간다. 아니라면 왼쪽의 $c$개를 건너뛰고 오른쪽에서 $k-c$번째 공을 찾는다. 이를 반복해 도착한 리프의 상자 번호가 답이다.

공 수가 음수가 아니고 $k$가 전체 공 수 이하이므로 항상 해당 공이 있는 리프에 도착한다. 조회는 위치만 찾을 뿐 공을 제거하지 않는다.

시간복잡도는 $O(N+Q\log(N+1))$, 공간복잡도는 $O(N)$이다. 공의 총수는 초기 개수와 추가된 개수의 합이며 $(N+Q)\cdot10^9$ 이하이다.

구현 · 언어별 풀이

C++ 구현

공의 총수와 $k$는 long long으로 둔다. 세그먼트 트리의 리프를 상자 번호순으로 배치하고, 선택 탐색에서 오른쪽으로 내려갈 때만 왼쪽 합을 $k$에서 뺀다. 조회에서는 트리의 값을 변경하지 않는다.

C 구현

세그먼트 트리의 원소 개수 합은 long long으로 저장한다. k도 같은 자료형으로 읽는다. 왼쪽 합 이상 여부로 한쪽 자식만 선택하고 오른쪽으로 가면 왼쪽 합을 k에서 뺀다.

Python 구현

합 트리를 array('q')로 두면 큰 상태의 객체 비용을 줄인다. k번째 위치는 root부터 한쪽으로만 내려가 찾으며, 질의 후 빈도를 삭제하지 않는다. 0인 빈도도 정상 원소 상태다.

Java 구현

long[] tree와 long k를 사용한다. 왼쪽 합이 k 이상이면 왼쪽, 아니면 k에서 왼쪽 합을 빼고 오른쪽으로 간다. 위치 찾기를 별도 이분 탐색과 구간합으로 반복하지 않는다.

Rust 구현

빈도 합과 k는 i64 또는 u64로, 노드 인덱스는 usize로 둔다. 한 단계마다 왼쪽 합과 비교해 내려간다. 부호 없는 k에서 뺄 때는 왼쪽 합보다 큰 분기에서만 수행한다.

JavaScript 구현

누적 공 개수는 최대 $2\cdot10^{15}$까지 커질 수 있다. Number의 정확한 정수 범위 안이지만 $32$비트에는 들어가지 않으므로 Fenwick 합은 Float64Array에 둔다. 인덱스와 인덱스의 최하위 비트만 비트 연산으로 처리한다.

$k$번째 공의 위치는 이분 탐색으로 prefix sum을 반복 계산하지 않고, Fenwick 트리에서 큰 $2$의 거듭제곱부터 내려가며 찾는다. 지금까지 건너뛴 공보다 더 필요하면 해당 구간 합을 $k$에서 빼고 인덱스를 옮긴다.

let pos = 0;
for (let step = topBit; step > 0; step >>= 1) {
    const next = pos + step;
    if (next <= n && tree[next] < k) {
        k -= tree[next];
        pos = next;
    }
}
const answer = pos + 1;

topBit는 $N$ 이하의 가장 큰 $2$의 거듭제곱이다. $k$는 $1$부터 세므로 비교는 < k이다. 답을 찾았다고 공을 제거하지 않는다. 초기 트리는 부모에게 합을 넘겨 선형 시간에 만든다.

제출

언어