P6041 「ACOI2020」布丁暗杀计划 题解

第100道蓝题是学长出的题欸。

---

可持久化线段树做法。

题目要求的其实是：在某一棵子树的同一深度中，$k$ 级祖先都相同，并且颜色和 $k$ 级祖先相同的结点的权值和。

我们依次来考虑这些条件。

首先是深度相同，最直接的想法是将深度相同的结点都放在一起，比如以深度为优先级对这些结点排序。

但是这样做的话会发现第二个条件就不容易处理了，不能这么直接简单粗暴地把深度相同的结点放在一起。

发现上面这种做法不可取的原因是同一棵子树中的结点排完序后可能不连续，其实只需要找到一种做法使得所有深度相同的结点是连续的并且在同一棵子树中的结点也是连续的，而广度优先搜索得到的搜索序正好符合这个条件。

于是我们先对这棵树 bfs 一下，得到结点的 bfs 序，并且 $k$ 级祖先相同的结点在 bfs 序中是连续的一段。

那怎么知道在这个序列中，$k$ 级祖先相同的一段区间的左端点和右端点呢？只需要分别二分出这个左端点和右端点即可。具体地，记录深度相同的结点在序列中的最左端和最右端，分别记为 $lt_{dep}$ 和 $rt_{dep}$，二分左端点的时候初始区间为 $[lt_{dep}, bfn_{now}]$，二分右端点的初始区间为 $[bfn_{now}, rt_{dep}]$，这里的 $now$ 表示待询问结点，$bfn$ 表示 bfs 序。

接下来就是最后一个问题了，求出在这一段中，颜色和 $k$ 级祖先相同的结点的权值和。

先把这个“颜色和 $k$ 级祖先相同”去掉，问题变成“区间查询和给定颜色相同的元素的权值和”。

如果只有一个询问的话可以开一个桶，从左端点开始遍历，过程中把遍历到的元素的权值加到它的颜色对应的桶里，知道遍历完这个区间。输出答案直接查询桶里的权值和即可。

变化到多组的话可以把桶变成前缀和，区间相减得到答案。但是这样的时空复杂度不能通过，用可持久化数组来维护即可，本质上就是可持久化线段树。

查询的时候判一下待查询结点有没有 $k$ 级祖先，检查深度即可。注意答案的范围是 $\sum\limits_{i = 1}^{n}d_{i} \leqslant 5 \times 10^{10}$，需要开 long long。

代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n, q, f, u, k, l, r, mid, posl, posr, father, color[500005], d[500005], dep[500005], lg[500005], bfn[500005], rnk[500005], lt[500005], rt[500005], fa[500005][20];
queue<int> qu;
vector<int> g[500005];
struct Segment_Tree {
    int idx, root[500005], lc[10000005], rc[10000005];
    ll sum[10000005];
    #define mid ((l + r) >> 1)
    int new_node(int k = 0) {
        ++idx;
        lc[idx] = lc[k], rc[idx] = rc[k], sum[idx] = sum[k];
        return idx;
    }
    int change(int k, const int& pos, const int& v, int l = 1, int r = 500000) {
        int rt = new_node(k);
        if(l == r) {
            sum[rt] += v;
            return rt;
        }
        if(pos <= mid) lc[rt] = change(lc[rt], pos, v, l, mid);
        else rc[rt] = change(rc[rt], pos, v, mid + 1, r);
        return rt;
    }
    ll ask(int kl, int kr, const int& pos, int l = 1, int r = 500000) {
        // cerr << pos << " " << l << " " << r << " " << sum[kl] << " " << sum[kr] << '\n';
        if(l == r) return sum[kr] - sum[kl];
        if(pos <= mid) return ask(lc[kl], lc[kr], pos, l, mid);
        else return ask(rc[kl], rc[kr], pos, mid + 1, r);
    }
    #undef mid
} tree;
int get(int x, int y) {
    while(y) {
        x = fa[x][lg[y]];
        y ^= (1 << lg[y]);
    }
    return x;
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    cin >> n >> q;
    for(int i = 1; i <= n; ++i) cin >> color[i];
    for(int i = 1; i <= n; ++i) cin >> d[i];
    for(int i = 2; i <= n; ++i) {
        cin >> f;
        g[f].push_back(i);
        lg[i] = lg[i >> 1] + 1;
    }
    memset(lt, 0x3f, sizeof(lt));
    qu.push(1);
    dep[1] = 1;
    for(int i = 1; i <= n; ++i) {
        u = qu.front();
        qu.pop();
        bfn[u] = i;
        rnk[i] = u;
        lt[dep[u]] = min(lt[dep[u]], i);
        rt[dep[u]] = max(rt[dep[u]], i);
        for(int j = 1; j <= lg[n]; ++j) fa[u][j] = fa[fa[u][j - 1]][j - 1];
        for(const auto& j : g[u]) {
            dep[j] = dep[u] + 1;
            fa[j][0] = u;
            qu.push(j);
        }
    }
    for(int i = 1; i <= n; ++i) {
        // cerr << rnk[i] << ": " << color[rnk[i]] << " " << d[rnk[i]] << '\n';
        tree.root[i] = tree.change(tree.root[i - 1], color[rnk[i]], d[rnk[i]]);
    }
    while(q--) {
        cin >> u >> k;
        if(dep[u] <= k) {
            cout << "0\n";
            continue;
        }
        father = get(u, k);
        l = lt[dep[u]], r = posl = bfn[u];
        while(l <= r) mid = (l + r) >> 1, (get(rnk[mid], k) == father) ? (posl = mid, r = mid - 1) : (l = mid + 1);
        l = posr = bfn[u], r = rt[dep[u]];
        while(l <= r) mid = (l + r) >> 1, (get(rnk[mid], k) == father) ? (posr = mid, l = mid + 1) : (r = mid - 1);
        // cerr << posl << " - " << posr << '\n';
        cout << tree.ask(tree.root[posl - 1], tree.root[posr], color[father]) << '\n';
    }
    return 0;
}