PT07J - Query on a tree III DFS序 + 主席树

Posted 2021-03-12 xxfy1

dfs序编号后跑权值主席树
但写起来是真的麻烦，总是wa，只能拿出模板过了

#include<bits/stdc++.h>
const int N = 100001;
using namespace std;

struct node {
    int to, ne;
} e[N<<1];

int n, m, len, l1, l2;
int id[N], rk[N], eid[N], re[N];
int he[N], a[N], b[N];
int val[N << 5], lc[N << 5], rc[N << 5], rt[N << 5];

void add(int x, int y) {
    e[++l2].to = y, e[l2].ne = he[x];
    he[x] = l2;
}

void dfs(int x, int fa) {
    id[x] = ++l1, rk[l1] = x;
    for(int i = he[x]; i; i = e[i].ne) {
        int y = e[i].to;
        if(y == fa) continue;
        dfs(y, x);
    }
    eid[x] = l1;
}

int change(int rt, int l, int r, int x, int y) {
    int now = ++len;
    lc[now] = lc[rt], rc[now] = rc[rt], val[now] = val[rt] + 1;
    if(l == r) {
        re[l] = y;
        return now;
    }
    int mid = (l + r) >> 1;
    if(x <= mid)
        lc[now] = change(lc[rt], l, mid, x, y);
    else
        rc[now] = change(rc[rt], mid + 1, r, x, y);
    return now;
}

int query(int u, int v, int l, int r, int k) {
    if(l == r)
        return re[l];
    int mid = (l + r) >> 1, s = val[lc[v]] - val[lc[u]], res;
    if(s >= k)
        res = query(lc[u], lc[v], l, mid, k);
    else
        res = query(rc[u], rc[v], mid + 1, r, k - s);
    return res;
}

int main() {
    scanf("%d", &n);
    for(int i = 1; i <= n; i++)
        scanf("%d", &a[i]), b[i] = a[i];
    sort(b + 1, b + n + 1);
    for( int i = 1; i < n; i++) {
        int x, y;
        scanf("%d%d", &x, &y);
        add(x, y);
        add(y, x);
    }
    dfs(1, 0);
    for(int i = 1; i <= n; i++) {
        a[rk[i]] = lower_bound(b + 1, b + n + 1, a[rk[i]]) - b;
        rt[i] = change(rt[i - 1], 1, n, a[rk[i]], rk[i]);
    }
    scanf("%d", &m);
    for(int i = 1; i <= m; i++) {
        int x, k;
        scanf("%d%d", &x, &k);
        printf("%d
", query(rt[id[x] - 1], rt[eid[x]], 1, n, k));
    }
    return 0;
}