백준 31603 Tree Quiz

https://www.acmicpc.net/problem/31603

 

클래스 8 문제 + 아챔 기출이다. 

 

문제는 몇 달 전에 읽었는데, (x-1)*n² + (LCA(x,y)-1)*n + (y-1)식 해석을 못 해서 이차식인가? 어쩌구.. 하다 접었었다. 잘 보면 x-1, lca(x,y)-1, y-1이 모두 0 이상 n-1 이하이니 n진수 숫자 하나로 생각할 수 있다. 즉, 각 쿼리는 (x, lca(x,y), y) 튜플 n^2개 중 k번째를 찾으라는 뜻이다. 

 

1, 2, ..., n으로 시작되는 튜플이 각각 정확히 n개씩 있으니 k번째 튜플의 x는 바로 확정되고, (x, -, -) 튜플 중 k₂ = (k-1) % N + 1 번째 튜플을 찾아야 한다. k번째를 찾을 때는 흔히 1..m까지 합이 k개 이상인지 여부로 이분탐색을 하니, 이분탐색이 가능하도록 잘 처리할 방법을 생각해 보자. lca(x,y) 의 후보는 x의 조상들이니, x의 임의의 조상 p에 대해 lca(x,y) <= p인 y의 개수를 알 수 있다면 이분탐색으로 lca(x,y)를 알 수 있다. 트리에서 조상의 정보가 필요한 경우에 쓸 수 있는 방법 중 하나는 dfs를 하며 조상의 정보를 추가/롤백하는 것이다. x, p가 고정되어 있을 때 lca(x,y) = p인 y의 개수는 (p의 서브트리 크기) - (p의 x 방향 자식의 서브트리 크기)와 같으니, 오프라인으로 쿼리를 처리하며 세그 등으로 지금 들어온 방향을 제외한 서브트리 크기를 잘 관리해 주면 이분탐색을 위한 쿼리를 $O(\log N)$에 처리할 수 있으니 튜플의 두 번째 항 lca(x,y)를 각 $O(\log^2 N)$에 구할 수 있다.

 

(x, lca(x,y), -) 튜플 중 t번째를 구해야 한다고 가정하자. x != lca(x,y)라면 lca(x,y)의 서브트리에는 있지만 x의 서브트리에는 없는 정점 중 t번째를 구해야 한다. ett를 사용하여 트리를 펴고 정점 번호를 수열의 값이라고 생각하면 일부 구간에서 t번째 수를 찾는 문제로 환원되고, 이는 잘 알려져 있다시피 pst 또는 머지 소트 트리를 사용하면 풀 수 있다. 머지소트트리를 쓰면 쿼리당 $O(\log^3 N)$, pst를 쓰면 쿼리당 (아마)$O(\log N)$까지도 될 것이다. 

 

pst를 바로 잘 짤 자신도 없었고 로그세제곱도 할 만 해 보여서 머지소트트리로 짰다. lca를 확정하면 lca가 더 작은 것들의 개수를 k₂에서 빼 t를 구해야 하는데, 이걸 까먹어서 디버깅 좀 했다. 그래도 틀리길래 #define int long long을 박았더니 잘 돌아갔다. 

 

#include <bits/stdc++.h>
#define int long long
using namespace std;
using ll = long long;
using pii = pair<ll,ll>;
 
ll sz[101010], in[101010], out[101010], id[101010], N, Q, cnt[101010], nxt[101010], ans[101010];
vector<ll> T[101010], mst[101010 << 2];
vector<pii> qry[101010];
 
void update(int i, int d) {
    for (; i <= N; i += i&-i)
        cnt[i] += d;
}
int query(int i) {
    int ret = 0;
    for (; i; i -= i&-i)
        ret += cnt[i];
    return ret;
}
 
void build(int v) {
    static int ii = 0;
    in[v] = ++ii;
    id[ii] = v;
    sz[v] = 1;
    for (int i : T[v]) {
        build(i);
        sz[v] += sz[i];
    }
    out[v] = ii;
}
 
void mst_build(int now, int l, int r) {
    if (l == r) {
        mst[now].push_back(id[l]);
        return;
    }
    int m = l+r >> 1;
    mst_build(now*2, l, m);
    mst_build(now*2+1, m+1, r);
    mst[now].resize(r-l+1);
    merge(mst[now*2].begin(), mst[now*2].end(), mst[now*2+1].begin(), mst[now*2+1].end(), mst[now].begin());
}
 
int mst_query(int now, int l, int r, int s, int e, int k) {
    if (r < s or e < l) return 0;
    if (s <= l and r <= e) return upper_bound(mst[now].begin(), mst[now].end(), k) - mst[now].begin();
    int m = l+r >> 1;
    return mst_query(now*2, l, m, s, e, k) + mst_query(now*2+1, m+1, r, s, e, k);
}
 
void dfs(int v) {
    update(v, sz[v]);
    for (auto[k,qi] : qry[v]) {
        int l = 0, r = N+1;
        while (l+1 < r) {
            int m = l+r >> 1;
            if (query(m) >= k) r = m;
            else l = m;
        }
        int lca = r;
        k -= query(r-1);
        l = 0, r = N+1;
        while (l+1 < r) {
            int m = l+r >> 1;
            if (mst_query(1,1,N,in[lca],out[lca],m) - mst_query(1,1,N,in[nxt[lca]],out[nxt[lca]],m) >= k)
                r = m;
            else
                l = m;
        }
        ans[qi] = (v-1)*N*N + (lca-1)*N + r-1;
    }
    for (int i : T[v]) {
        update(v, -sz[i]);
        nxt[v] = i;
        dfs(i);
        update(v, sz[i]);
    }
    update(v, -sz[v]);
}
 
signed main() {
    cin.tie(nullptr)->sync_with_stdio(0);
 
    cin >> N >> Q;
    int root = 0;
    for (int i = 1, p; i <= N; i++){
        cin >> p;
        if (p) T[p].push_back(i);
        else root = i;
    }
    build(root);
    mst_build(1,1,N);
    for (int i = 1,x; i <= Q; i++) {
        cin >> x;
        qry[(x-1)/N+1].emplace_back((x-1) % N + 1, i);
    }
    dfs(root);
    for (int i = 1; i <= Q; i++) cout << ans[i] << '\n';
}