백준 14898번 서로 다른 수와 쿼리 2

문제 요약

길이 N인 배열이 주어진다. l r로 쿼리가 주어진다. 구간 [l,r]에서 서로 다른 수의 개수를 출력해야 한다.

N은 최대 100만이다.

풀이

숫자가 크다. 좌표압축부터 하고 시작하자.

 

[1...x]까지를 나타내는 $segtree_x$를 생각해보자.

이 세그트리에서 1이 켜지는 위치는 해당 구간에서 한 번만 등장하는 수는 그 수가 등장하는 위치에 1이 켜지고, 여러 번 등장하는 수가 있다면 그 수가 등장하는 곳들 중에서 제일 오른쪽에만 1이 켜진다.

예를 들어서 배열이 [1,1,4,2,1,2]라고 하면 [0,0,1,0,1,1]이 된다.

 

이런 세그먼트 트리가 있다면, [l,r] 쿼리가 들어오면 $segtree_r$에서 [l,r]로 구간합 쿼리를 날리면 우리가 원하는 답이 된다.

이런 세그먼트 트리를 만들기 위해서 $prv[i]$라는 값을 먼저 전처리해둘 필요가 있다. $prv[i]$란 $i$번째 수가 이전에 등장한 위치이다. $i$번째 수가 처음 등장하는 거라면 $prv[i]$는 -1이다.

 

제일 처음엔 모든 수가 0인 세그트리를 만든다 이를 $segtree_0$이라고 하자.

$segtree_{i+1}$을 만들기 위해서는 $segtree_i$에서 $i+1$의 위치에 1을 증가시킨다. 그 다음에 $prv[i+1]$이 -1이 아니라면 이전에 등장했다는 뜻이므로 $prv[i+1]$의 위치에 1 을 감소시킨다. 그 결과가 $segtree_{i+1}$이 된다.

 

이제 $segtree_{[0...N]}$을 모두 들고 있어야하는데 이는 PST로 해결이 가능하다.

#include<bits/stdc++.h>
using namespace std;
const int MXN = 1e6;

int a[MXN+5];
int start_num[MXN+5], end_num[MXN+5];
vector<int> v, roots(MXN+5);
int N, Q;
struct Node {
    int l, r;
    int cnt;
    Node() : cnt(0) {};
};
vector<Node> tree(MXN*50);
int node_cnt;
void init() {
    node_cnt = N << 1;
    for(int i=1;i<N;++i) {
        tree[i].l = i << 1;
        tree[i].r = i << 1 | 1;
    }
}

void update(int node, int s, int e, int idx, int val) {
    tree[node].cnt += val;
    if(s != e) {
        int mid = s+e >> 1;
        int n1 = tree[node].l, n2 = tree[node].r;
        if(idx <= mid) {
            tree[node].l = node_cnt;
            tree[node_cnt++] = tree[n1];
            update(tree[node].l, s, mid, idx, val);
        }
        else {
            tree[node].r = node_cnt;
            tree[node_cnt++] = tree[n2];
            update(tree[node].r, mid+1, e, idx, val);
        }
    }
}

int query(int node, int s, int e, int l, int r) {
    if(l <= s && e <= r) return tree[node].cnt;
    if(e < l || s > r) return 0;
    int mid = s + e >> 1;
    int n1 = tree[node].l, n2 = tree[node].r;
    return query(n1, s, mid, l, r) + query(n2, mid+1, e, l, r);
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N;
    for(int i=1;i<=N;++i) {
        cin >> a[i]; v.push_back(a[i]);
    }
    sort(v.begin(), v.end());
    v.erase(unique(v.begin(), v.end()), v.end());
    for(int i=1;i<=N;++i) a[i] = lower_bound(v.begin(), v.end(), a[i]) - v.begin();
    vector<int> prev_idx(v.size(), 0);
    vector<int> prv(N+1, -1);
    for(int i=1;i<=N;++i) {
        if(prev_idx[a[i]] != 0) {
            prv[i] = prev_idx[a[i]];
        }
        prev_idx[a[i]] = i;
    }
    init();
    roots[0] = 1;
    int cur_root = roots[0];
    for(int i=1;i<=N;++i) {
        tree[node_cnt] = tree[cur_root];
        cur_root = node_cnt++;
        update(cur_root, 1, N, i, 1);
        if(prv[i] != -1) {
            tree[node_cnt] = tree[cur_root];
            cur_root = node_cnt++;
            update(cur_root, 1, N, prv[i], -1);
        }
        roots[i] = cur_root;
    }
    cin >> Q;
    int ans = 0;
    for(int i=0;i<Q;++i) {
        int x, r; cin >> x >> r;
        int l = x + ans;
        assert(l <= r);
        ans = query(roots[r], 1, N, l, r);
        cout << ans << '\n';
    }
    return 0;
}