백준 11932번 트리와 K번째 수

문제 요약

크기가 N인 트리가 주어진다. 각 노드에는 숫자가 있다.

X Y K 로 쿼리가 주어지는데 이는 X번 노드부터 Y번 노드까지 가는 경로 중에서 K번째 수를 출력하는 것이다.

트리의 크기와 쿼리 수 둘 다 최대 10만이다.

풀이

선생님 이거 문제가 너무 어렵지 않습니까

일단 트리에서 경로를 빠르게 찾아야 한다. 만약에 어떤 자료구조가 루트부터 시작해서 $i$번 정점까지 가는 경로를 나타내고 있다고 하자.

 

그런 자료구조를 $d_i$라고 하면 X부터 Y까지 가는 경로는 $d_X + d_Y - d_{LCA(X,Y)} - d_{par(LCA(x,y))}$라고 할 수 있다.

무슨 소린지 모르겠어도 괜찮다.

 

일단 정점에 적혀 있는 숫자들을 전부 좌표압축하고 압축한 수로 바꿔주자.

그럼 대략 0부터 N-1까지라고 할 수 있다.

이제 이런 세그트리를 생각하자. $segtree_x$는 루트(1번)노드 부터 $x$번 노드까지 가는 경로에 있는 수들의 위치에 1씩 더해준 세그트리다.

이제 이런 트리에서 $K$번째 수를 찾는다면 이것은 루트부터 $x$번 노드까지 가는 경로에 있는 수들 중에서 $K$번째 수를 찾는 것과 같다.

이 $segtree_x$는 저 위에서 말한 $d$로 쓴 식으로 난리 친것과 같아졌다.

 

이제 세그트리를 N개 만들어야 된다는 힘겨운 난관만 지나면 된다.

$segtree_{par}$을 어떤 노드 $x$의 부모 노드의 세그트리라고 하자. 그러면 $segtree_x$는 $segtree_{par}$과 비교했을 때 딱 1만 차이난다.

 

즉, 처음에 전부 0인 세그트리를 하나 만들고 DFS를 진행하면서 현재 노드의 세그트리는 부모 노드의 세그트리에서 한 부분만 업데이트를 진행해주면 된다.

이는 PST로 해결이 가능하고 메모리 문제도 해결할 수 있다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const int MAXN = 1e5+5;
int N, M;
int depth[MAXN], sp_table[20][MAXN];
vector<vector<int>> G(MAXN);
vector<int> v;
void dfs(int cur, int par, int d) {
    depth[cur] = d;
    for(const int &v:G[cur]) {
        if(v == par) continue;
        dfs(v, cur, d+1);
        sp_table[0][v] = cur;
    }
}

int LCA(int u, int v) {
    if(depth[u] > depth[v]) swap(u,v);
    int diff = depth[v] - depth[u];
    int idx = 0;
    while(diff) {
        if(diff & 1) v = sp_table[idx][v];
        diff >>= 1; ++idx;
    }
    if(u == v) return u;
    for(int i=19;i>=0;--i) {
        if(sp_table[i][u] != sp_table[i][v]) {
            u = sp_table[i][u];
            v = sp_table[i][v];
        }
    }
    return sp_table[0][u];
}

struct Node {
    int l, r, val;
    Node() : l(0), r(0), val(0) {};
    Node(int l_, int r_, int val_) : l(l_), r(r_), val(val_) {};
};

vector<int> a, indices;
vector<Node> tree(200005);
vector<int> roots(100005);
int parent[100005];

void init() {
    for(int i=1;i<v.size();++i) {
        tree[i].l = i << 1;
        tree[i].r = i << 1 | 1;
    }
}

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

void make_tree(int cur, int par) {
    parent[cur] = par;
    roots[cur] = tree.size(); tree.push_back(tree[roots[par]]);
    update(roots[cur], 0, v.size()-1, 1, indices[cur]);
    for(auto nxt:G[cur]) {
        if(nxt == par) continue;
        make_tree(nxt, cur);
    }
}

int query(int s, int e, int lca, int p, int k) {
    int l = 0, r = v.size() - 1;
    while(l != r) {
        int mid = l+r >> 1;
        int lsz = tree[tree[e].l].val + tree[tree[s].l].val - tree[tree[lca].l].val - tree[tree[p].l].val;
        if(lsz >= k) {
            s = tree[s].l; e = tree[e].l;
            lca = tree[lca].l; p = tree[p].l;
            r = mid;
        }
        else {
            k -= lsz;
            s = tree[s].r; e = tree[e].r;
            lca = tree[lca].r; p = tree[p].r;
            l = mid+1;
        }
    }
    return l;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> M;
    a.resize(N+1);
    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());
    indices.resize(N+1);
    for(int i=1;i<=N;++i) indices[i] = lower_bound(v.begin(), v.end(), a[i]) - v.begin();
    for(int i=1;i<N;++i) {
        int u, v; cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, 0, 0);
    for(int i=1;i<20;++i) for(int j=1;j<=N;++j) sp_table[i][j] = sp_table[i-1][sp_table[i-1][j]];
    init();
    roots[0] = 1;
    make_tree(1, 0);
    for(int i=0;i<M;++i) {
        int x, y, k; cin >> x >> y >> k;
        if(x == y) {
            cout << a[x] << '\n';
            continue;
        }
        int l = LCA(x, y);
        cout << v[query(roots[x], roots[y], roots[l], roots[parent[l]], k)] << '\n';
    }
    return 0;
}