Ordered Set

GNU 확장인 PBDS의 순서 통계 트리를 사용한다. 표준 C++ 라이브러리는 아니며 GNU C++ 환경에서 다음 헤더와 타입을 사용한다.

#include <functional>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using ordered_set = __gnu_pbds::tree<
    int, __gnu_pbds::null_type, std::less<int>,
    __gnu_pbds::rb_tree_tag,
    __gnu_pbds::tree_order_statistics_node_update>;

ordered_set s에서 s.order_of_key(x)는 $x$보다 작은 원소의 수이다. s.find_by_order(i)는 $0$부터 센 인덱스 $i$의 원소를 가리키는 반복자를 반환하므로, 값은 *s.find_by_order(i)로 읽는다. 이 문제에서는 인덱스가 항상 유효하다. 삽입·삭제·크기는 insert(x), erase(x), size()로 처리한다.

비교자는 엄격한 비교인 std::less<int>로 두어 중복을 저장하지 않는다. less_equal로 바꾸어 중복 집합을 흉내 내면 트리의 비교자 계약을 깨뜨린다. 일반 set에서 반복자를 $i$번 이동하는 것은 로그 시간이 아니므로 순위 질의의 대체 구현으로 사용하지 않는다.

C 표준 라이브러리에는 순위 질의를 지원하는 집합이 없다. 균형 트리를 직접 만드는 대신 좌표 압축과 Fenwick tree를 사용한다.

질의를 먼저 저장하고 삽입되는 값을 정렬·중복 제거하여 values에 둔다. 압축 인덱스는 $0$부터, Fenwick tree의 인덱스는 $1$부터 센다. 존재 여부가 바뀌는 삽입·삭제에서만 해당 위치의 빈도를 $+1$ 또는 $-1$ 갱신한다.

$x$보다 작은 원소의 수는 values에서 $x$의 lower bound를 구한 뒤, 그 위치 앞까지의 누적합으로 얻는다. 빈도가 $0$ 또는 $1$이므로 누적합은 실제 집합의 원소 수와 같다.

삭제할 값의 lower bound를 j라 하면 j < n && values[j] == x && present[j]인 경우에만 존재 표식을 지우고 add(j+1, -1)을 호출한다. 삽입 값만 압축했으므로, 삽입된 적 없는 값의 삭제에서는 다음 값의 빈도를 잘못 줄이지 않아야 한다.

int n; int *bit, *values;를 전역 변수로 두고, bit은 calloc(n+1, sizeof *bit)로 $0$으로 초기화한다. values에는 압축된 원래 값을 저장한다. 아래 add에는 $1$부터 센 위치, prefix에는 포함할 압축 원소의 개수를 전달한다.

void add(int i, int delta) {
    for (; i <= n; i += i & -i) bit[i] += delta;
}

int prefix(int i) {
    int sum = 0;
    for (; i > 0; i -= i & -i) sum += bit[i];
    return sum;
}

int kth(int k) {
    int idx = 0, step = 1;
    while (step <= n / 2) step <<= 1;
    for (; step; step >>= 1) {
        int next = idx + step;
        if (next <= n && bit[next] <= k) {
            k -= bit[next];
            idx = next;
        }
    }
    return values[idx];
}

여기서 k는 $0$부터 센 유효한 순위다. 현재 위치 다음 구간의 원소 수가 남은 순위 k 이하이면 그 구간을 건너뛰고, 건너뛴 원소 수만큼 k를 줄인다. 큰 $2$의 거듭제곱부터 확인하면 마지막 idx는 원하는 값의 압축 인덱스가 된다. 원래 값은 values[idx]로 읽는다.

좌표 압축은 값의 대소 관계를 보존한다. 삽입·삭제는 존재 여부를 정확히 반영하고, 누적합과 순위 탐색은 현재 원소 수를 기준으로 하므로 모든 질의가 원래 집합과 같은 답을 낸다. 정렬과 각 질의를 합쳐 시간복잡도는 $O(Q\log Q)$, 공간복잡도는 $O(Q)$이다.

값의 존재 여부를 누적하기

Python 표준 라이브러리에는 순서 통계 트리가 없다. 정렬된 리스트의 삽입·삭제는 원소를 옮기는 데 선형 시간이 걸리므로, 모든 질의를 먼저 읽고 삽입될 수 있는 값만 정렬하여 좌표를 압축한다. 압축은 값의 대소 관계를 유지한다.

압축한 각 위치에 현재 집합에 있으면 $1$, 없으면 $0$을 저장하고 Fenwick Tree로 누적합을 관리한다. 이미 있는 값의 삽입과 없는 값의 삭제는 무시한다. $x$보다 작은 압축 좌표의 개수를 bisect_left로 구하면, 그 앞까지의 누적합이 order_of_key x의 답이다.

find_by_order i에서는 누적합이 $i+1$ 이상인 첫 위치를 찾는다. Fenwick Tree의 이진 리프팅으로 큰 구간부터 탐색하여 누적합이 $i+1$ 미만인 접두 구간을 늘린다. 탐색이 끝나면 그 다음 원소의 압축 좌표가 답이다. 존재 여부가 $0$ 또는 $1$이므로 누적합이 증가하는 순서가 집합의 오름차순과 같다.

좌표 개수를 $K$라고 하면 정렬은 $O(Q\log(Q+1))$, 각 갱신·질의는 $O(\log(K+1))$이다. 전체 시간복잡도는 $O(Q\log(Q+1))$, 공간복잡도는 질의 저장을 포함하여 $O(Q)$이다.

import sys
from bisect import bisect_left

read = sys.stdin.buffer.readline
queries = [read().split() for _ in range(int(read()))]
values = sorted({int(q[1]) for q in queries if q[0] == b"insert"})
position = {x: j for j, x in enumerate(values)}
n = len(values)
tree = [0] * (n + 1)
present = [False] * n
size = 0
answers = []

def add(j, delta):
    j += 1
    while j <= n:
        tree[j] += delta
        j += j & -j

def prefix(j):
    result = 0
    while j:
        result += tree[j]
        j -= j & -j
    return result

def kth(order):
    j = 0
    step = 1 << (n.bit_length() - 1)
    while step:
        nxt = j + step
        if nxt <= n and tree[nxt] <= order:
            order -= tree[nxt]
            j = nxt
        step >>= 1
    return values[j]

for q in queries:
    command = q[0]
    if command == b"size":
        answers.append(str(size))
        continue
    x = int(q[1])
    if command == b"insert":
        j = position[x]
        if not present[j]:
            present[j] = True
            add(j, 1)
            size += 1
    elif command == b"erase":
        j = position.get(x)
        if j is not None and present[j]:
            present[j] = False
            add(j, -1)
            size -= 1
    elif command == b"order_of_key":
        answers.append(str(prefix(bisect_left(values, x))))
    else:
        answers.append(str(kth(x)))
sys.stdout.write("\n".join(answers))

kth를 호출하는 질의는 유효한 인덱스만 주어지므로 그때는 $K>0$이다. 삽입 질의가 하나도 없더라도 크기·순위 질의는 그대로 처리된다.

TreeSet은 정렬 순서를 유지하지만 빠른 순위 조회는 제공하지 않는다. 반복자를 이동하면 선형 시간이 되므로 좌표 압축과 Fenwick tree로 순위 질의를 처리한다.

질의를 먼저 저장하고 삽입되는 값을 정렬·중복 제거하여 values에 둔다. 압축 인덱스는 $0$부터, Fenwick tree의 인덱스는 $1$부터 센다. 존재 여부가 바뀌는 삽입·삭제에서만 해당 위치의 빈도를 $+1$ 또는 $-1$ 갱신한다.

$x$보다 작은 원소의 수는 values에서 $x$의 lower bound를 구한 뒤, 그 위치 앞까지의 누적합으로 얻는다. 빈도가 $0$ 또는 $1$이므로 누적합은 실제 집합의 원소 수와 같다.

삭제할 값의 lower bound를 j라 하면 j < n && values[j] == x && present[j]인 경우에만 존재 표식을 지우고 add(j+1, -1)을 호출한다. 삽입 값만 압축했으므로, 삽입된 적 없는 값의 삭제에서는 다음 값의 빈도를 잘못 줄이지 않아야 한다.

static int[] bit, values와 압축된 값의 개수 n을 두고 다음 함수를 사용한다. add에는 $1$부터 센 위치, prefix에는 포함할 압축 원소의 개수를 전달한다. 빈도와 누적합은 최대 질의 수이므로 int로 충분하다.

static void add(int i, int delta) {
    for (; i <= n; i += i & -i) bit[i] += delta;
}

static int prefix(int i) {
    int sum = 0;
    for (; i > 0; i -= i & -i) sum += bit[i];
    return sum;
}

static int kth(int k) {
    int idx = 0;
    for (int step = Integer.highestOneBit(n); step > 0; step >>= 1) {
        int next = idx + step;
        if (next <= n && bit[next] <= k) {
            k -= bit[next];
            idx = next;
        }
    }
    return values[idx];
}

여기서 k는 $0$부터 센 유효한 순위다. 현재 위치 다음 구간의 원소 수가 남은 순위 k 이하이면 그 구간을 건너뛰고, 건너뛴 원소 수만큼 k를 줄인다. 큰 $2$의 거듭제곱부터 확인하면 마지막 idx는 원하는 값의 압축 인덱스가 된다. 원래 값은 values[idx]로 읽는다.

좌표 압축은 값의 대소 관계를 보존한다. 삽입·삭제는 존재 여부를 정확히 반영하고, 누적합과 순위 탐색은 현재 원소 수를 기준으로 하므로 모든 질의가 원래 집합과 같은 답을 낸다. 정렬과 각 질의를 합쳐 시간복잡도는 $O(Q\log Q)$, 공간복잡도는 $O(Q)$이다.

BTreeSet::iter().nth(k)는 원소를 $k$번 이동하므로 순위 질의에 선형 시간이 걸린다. 좌표 압축과 Fenwick tree를 사용하고 Vec<bool>에 존재 여부, Vec<usize>에 누적 개수를 저장한다.

질의를 먼저 저장하고 삽입되는 값을 sort_unstable(), dedup()하여 values에 둔다. 압축 인덱스는 $0$부터, Fenwick tree의 인덱스는 $1$부터 센다. 존재 여부가 바뀌는 삽입·삭제에서만 해당 위치의 빈도를 $+1$ 또는 $-1$ 갱신한다.

$x$보다 작은 원소의 수는 values에서 $x$의 lower bound를 구한 뒤, 그 위치 앞까지의 누적합으로 얻는다. 빈도가 $0$ 또는 $1$이므로 누적합은 실제 집합의 원소 수와 같다.

lower bound는 values.partition_point(|&v| v < x)로 구한다. 아래 add에는 $1$부터 센 위치를 전달하고, prefix에는 포함할 압축 원소의 개수를 전달한다. 삭제는 존재하는 원소에서만 호출하므로 빈도를 뺄 때 언더플로하지 않는다.

삭제에서는 찾은 위치 j가 j < values.len() && values[j] == x && present[j]를 만족할 때만 존재 표식을 지우고 add(&mut bit, j+1, false)를 호출한다. 삽입된 적 없는 값은 압축 배열에 없을 수 있으므로 범위와 값의 일치를 먼저 확인한다.

fn add(bit: &mut [usize], mut i: usize, insert: bool) {
    while i < bit.len() {
        if insert { bit[i] += 1; } else { bit[i] -= 1; }
        i += i & i.wrapping_neg();
    }
}

fn prefix(bit: &[usize], mut i: usize) -> usize {
    let mut sum = 0;
    while i > 0 {
        sum += bit[i];
        i &= i - 1;
    }
    sum
}

fn kth(bit: &[usize], values: &[i32], mut k: usize) -> i32 {
    let n = values.len();
    let mut idx = 0;
    let mut step = 1;
    while step <= n / 2 { step <<= 1; }
    while step > 0 {
        let next = idx + step;
        if next <= n && bit[next] <= k {
            k -= bit[next];
            idx = next;
        }
        step >>= 1;
    }
    values[idx]
}

여기서 k는 $0$부터 센 유효한 순위다. 현재 위치 다음 구간의 원소 수가 남은 순위 k 이하이면 그 구간을 건너뛰고, 건너뛴 원소 수만큼 k를 줄인다. 큰 $2$의 거듭제곱부터 확인하면 마지막 idx는 원하는 값의 압축 인덱스가 된다. 원래 값은 values[idx]로 읽는다.

좌표 압축은 값의 대소 관계를 보존한다. 삽입·삭제는 존재 여부를 정확히 반영하고, 누적합과 순위 탐색은 현재 원소 수를 기준으로 하므로 모든 질의가 원래 집합과 같은 답을 낸다. 정렬과 각 질의를 합쳐 시간복잡도는 $O(Q\log Q)$, 공간복잡도는 $O(Q)$이다.

JavaScript의 Set은 삽입 순서로 순회하며 값의 순위나 $i$번째 값을 로그 시간에 구해 주지 않는다. 매번 정렬하기보다 질의를 먼저 읽고 삽입될 수 있는 값들을 숫자순으로 정렬·중복 제거하여 좌표를 압축한다.

Map으로 값에서 좌표로 가는 관계를 만들고, Uint8Array에 존재 여부, Int32Array Fenwick Tree에 각 좌표의 $0/1$을 저장한다. 중복 삽입과 없는 값의 삭제는 무시한다. order_of_key x는 압축 배열에서 $x$ 이상인 첫 위치를 이분 탐색한 뒤 그 앞의 개수 합을 구한다.

find_by_order i는 누적합이 $i+1$이 되는 첫 좌표를 Fenwick Tree 이진 리프팅으로 찾는다. 큰 $2$의 거듭제곱부터, 후보 접두합이 아직 $i+1$ 미만이면 그 구간을 소비한다. 존재 여부가 $0/1$이므로 누적 개수의 순서가 집합의 값 순서와 같다.

function kth(order) {
    let j = 0, step = 1;
    while (step * 2 <= n) step *= 2;
    for (; step > 0; step = Math.floor(step / 2)) {
        const next = j + step;
        if (next <= n && tree[next] <= order) {
            order -= tree[next];
            j = next;
        }
    }
    return values[j];
}

위 함수의 tree는 1-based, values는 0-based이며 order도 0-based이다. 유효한 순위 질의가 있을 때 $n>0$이다. 좌표 압축 정렬과 각 연산의 로그 탐색을 합쳐 전체 $O(Q\log(Q+1))$ 시간, 질의 저장을 포함한 $O(Q)$ 공간이다.