원소 추가, 삭제, 중간값 찾기 쿼리

2020 SCPC 2차 예선 3번 문제에서 삽질한 기록을 정리하려고 한다.

문제는 다음과 같다.

먼저 N<=200,000 개의 초기값이 주어진다. 그 후 Q<=1,000 번의 쿼리가 주어지는데 쿼리 하나는 순서대로 K<=400 번의 기존 원소 삭제, K번의 원소 추가, 1번의 중간값 찾기로 이루어져 있다 (즉 원소 추가는 총 K*Q번, 원소 삭제는 총 K*Q번, 중간값 계산은 Q번 주어진다). 주어지는 원소는 모두 1 이상 M<=20,000,000 이하의 자연수이고, 원소는 중복될 수 있으며 중간값 계산에는 중복 원소를 전부 포함시킨다. 중간값의 정의는 원소가 총 N개일 때 0-base로 floor(N/2)번째 값을 뜻한다. 즉 1, 2, 7, 9가 있으면 7이 중간값이다.

https://www.acmicpc.net/problem/1655 이 문제와 유사한데, 원소의 삭제가 이루어진다는 점과 중간값 계산의 빈도가 다르다.

1655번: 가운데를 말해요

일단 N이 20만이고, 추가되는 원소의 수와 삭제되는 원소의 수는 항상 동일하므로 한 쿼리가 끝날 때마다 전체 원소의 수는 일정하게 유지된다. 따라서 원소 삭제/추가/중간값 찾기를 각각 O(logN)에 수행하면 전체 시간복잡도는 O(KQlogN)이 되니까 괜찮은 솔루션이 될 것이다. 또는 원소 추가/삭제가 중간값 계산보다 빈번하게 일어난다는 점을 감안해, 추가/삭제의 시간복잡도를 좀 더 낮추고 중간값 계산의 시간복잡도를 어느정도 높여도 문제를 풀 수 있을 것이다.

이제 가능한 구현 4가지 정도를 설명하겠다. 각 구현은 시간복잡도가 다를 수도 있고, 시간복잡도는 같은데 상수가(실제 코드가 얼마나 빠른지) 다를 수도 있으며, 원소 값의 범위에 의존성이 있는 등 추가적인 제약 사항이 있을 수도 있다.

1. 주어지는 원소를 min heap과 max heap에 반반씩 넣어서 median을 구하기

아마 중간값 찾기 문제에서 가장 유명한 해법일 것이다.

모든 원소를 크기에 따라 절반으로 나눠서 작은 쪽은 max heap에, 큰 쪽은 min heap에 넣어두면 min heap의 root 원소가 중간값이 된다.
원소 추가/삭제를 heap의 push로 관리하고, min heap / max heap 간의 순서를 잘 보장해주는 작업도 상수번의 heap push / pop으로 관리할 수 있다.

문제는 C++ STL에서 제공하는 heap의 구현체(priority_queue)의 경우 내부 원소의 삭제를 지원하지 않는다. 가능한 해법은 각 원소의 인덱스를 유지해서 원소 삭제를 지원하는 heap을 직접 구현하거나, 메모리를 조금 희생하더라도 삭제할 원소 목록을 유지하다가 삭제되었어야 하는 원소가 heap의 root가 되었을 때 제거하는 방식이 있다.

후자의 방식은 최악의 경우 모든 삭제 대상 원소가 heap의 root가 되지 않아 추가 원소를 전부 다 메모리에 올리는 경우가 생길 수 있는데, SCPC에서는 추가되는 원소의 수 K*Q가 기존 원소의 수 N의 2배 정도이기 때문에 삭제할 원소 목록만 효율적으로 관리하면 메모리 제한에는 걸리지 않는다고 한다. 본인은 unordered_map으로 삭제 대상 원소를 관리했는데 메모리 부족으로 인해 틀렸고, 대신 또 다른 힙을 만들어서 삭제 대상 원소만 관리한 친구는 맞았다고 한다.

2. (1)의 해법에서 heap 대신 balanced binary search tree(STL의 set 등)를 사용하기

BBST의 경우 heap과 마찬가지로 총 원소의 개수가 #일 때 원소 추가를 O(log#)에 지원하며, 원소 삭제 또한 O(log#)으로 지원하고 STL(set)에도 구현되어 있다. 따라서 원소 삭제를 lazy하게 하지 않아도 된다는 장점이 있지만, STL의 set은 STL priority_queue와 비교했을 때 시간복잡도는 같아도 오버헤드(상수시간)가 매우 크기 때문에 해당 방법으로 접근하면 시간초과가 난다.

STL의 set은 해당 트리 안의 k번째 원소가 무엇인지 찾는 기능이 없기 때문에 (1)의 해법과 같이 2개의 set을 써야 하지만, g++를 사용한다면 k번째 원소 찾기를 지원하는 set을 사용할 수 있다. ( https://www.facebook.com/tamreffingbaekjoon/posts/467225557034913/ 참조 )
이걸 쓰면 그냥 모든 원소를 아예 set에 넣고 빼고를 반복하는 단순한 구현이 가능한데, 그래도 시간초과가 난다는 사실은 변함이 없다.

3. [1, M] 구간을 한 버킷의 크기가 sqrt(M)인 sqrt(M)개의 버킷으로 분할하기

만약 우리가 각 원소의 등장 빈도를 기록한 크기가 M인 배열과, 그의 누적합을 저장한 배열 C를 가지고 있다면 중간값은 C[M] / 2 <= C[x] 가 되는 최소의 x 값과 같게 된다. 따라서 원소가 바뀔 때 누적합 갱신과, 주어진 누적합 배열에서 x를 찾아내는 걸 빠르게 할 수 있다면 원 문제를 풀 수 있다.

이제 원소의 등장 빈도를 버킷 단위로 저장해보자. 한 버킷은 크기가 sqrt(M)인 배열이고, 이런 버킷을 sqrt(M)개 만들면 저장할 수 있는 값의 범위는 sqrt(M)^2 == M이 된다. 원소 v의 등장 빈도를 알고 싶다면 v / sqrt(M) 번째 버킷의 v % sqrt(M) 번째 배열값을 참조하면 된다.

그럼 원 문제의 요구사항을 다음과 같이 처리할 수 있다

원소 v 추가 == 대응되는 버킷 위치의 값을 1 증가, 해당 버킷 내 등장빈도 총 합 1 증가 => O(1)
원소 v 삭제 == 대응되는 버킷 위치의 값을 1 감소, 해당 버킷 내 등장빈도 총 합 1 감소 => O(1)
중간값 찾기 == "버킷별 등장빈도의 합" 을 저장한 배열을 순차적으로 순회하면서 그 누적합이 C[M] / 2 이상이 되는 첫번째 버킷을 찾고, 그 버킷 안의 정확히 몇번째 원소에서 누적합이 목표값을 넘는지를 찾는다. 버킷은 sqrt(M)개 있고, 한 버킷 내 원소도 sqrt(M)개 있기 때문에 시간복잡도는 O(sqrt(M) + sqrt(M)) = O(sqrt(M))

총 시간복잡도는 원소 추가/삭제 = O(KQ), 중간값 찾기 = O(Qsqrt(M)) 이 된다.
실제 SCPC에서는 해당 구현으로 맞았다. 중간값을 찾는 빈도가 원소 추가/삭제보다 매우 작았기 때문에 O(KQlogN)보다 O(KQ + Qsqrt(M))이 더 빨랐기 때문일 것이다.

4. [1, M] 를 index로 두고 해당 index의 원소가 몇번이나 등장했는지를 value로 가지는 세그먼트 트리 or BIT 만들기

즉 구간합을 저장하는 세그먼트 트리를 만들어서 사용하는 것이다. 기존 문제의 요구사항은 다음과 같이 세그먼트 트리 쿼리로 바꿀 수 있다.

원소 v 추가 == 세그먼트 트리 v번 위치의 값을 1 증가 (점쿼리)
원소 v 감소 == 세그먼트 트리 v번 위치의 값을 1 감소 (점쿼리)
중간값 찾기 == 세그먼트 트리의 root부터 출발해서, [1, x]의 누적합이 [1, M]의 누적합의 정확히 절반이 되는 최소의 x 를 찾아 한층씩 내려감 (일종의 binary search)

각 세그먼트 트리 쿼리의 시간복잡도는 트리의 높이에 비례하기 때문에 모두 O(logM) 의 시간이 들고, 총 시간복잡도는 O(KQlogM)이 된다.

 

SCPC의 경우 시간이 지나면 과거 문제를 채점할 수 있게 해주므로 채점이 가능해지면 코드도 추가로 올릴 예정이다.