P6041 「ACOI2020」布丁暗杀计划 Solution

· · 题解

题解区貌似没有我大 ds 线段树合并的做法，来交一发。

以下记 d_i 表示结点 i 在 原树 上的深度。

首先考虑如果颜色只有一种该怎么做，类似 CF208E Blood Cousins，首先对询问处理，先跑一个倍增或者长链剖分，把询问挂到 k 级祖先上，然后变成了查询一个点的 k 级儿子权值和。考虑对每个节点维护一个数组 a_{i,j}，表示在结点 i 的子树内原树深度为 j 的结点的权值和，那么对于一个询问 k，答案就是 a_{i,d_i+k}。维护这个数组就是先遍历子节点，子节点搜完之后合并到自己的数组上，最后把自己的权值加上去。然后这个合并可以用动态开点线段树合并实现，时间复杂度 O(n\log n)。

然后考虑分开颜色。我们发现不同的颜色之间互不影响，所以我们对每一种颜色建一个虚树，然后和上面一样的方法跑，因为总点数还是 n 不变，所以时间复杂度还是 O(n\log n)。

建树的过程就是在 DFS 的过程中维护一个 l_i 表示最近的颜色为 i 的祖先，然后每个点向与他相同颜色的最近祖先连边，然后就建完虚树了。

注意一个细节是这样可能不连通所以要每个颜色建一个虚点，作为最开始的 l_i。

#include <bits/stdc++.h>
#define LL long long
using namespace std;
const int N = 1e6 + 10;
int n, Q, col[N], W[N], Ans[N]; bool has[N]; vector<int> G[N];

// Rebuild and initial.
int depth[N], lst[N], fa[22][N]; vector<int> F[N];
void DFS(int u, int f) {
    int tmp = lst[col[u]]; lst[col[u]] = u; fa[0][u] = f; depth[u] = depth[f] + 1;
    if (tmp) F[tmp].emplace_back(u); else F[n + col[u]].emplace_back(u); 
    for (int v : G[u]) { DFS(v, u); } lst[col[u]] = tmp; return ;
}
vector<pair<int, int> > Qry[N];

namespace Sgt {

const int M = 1e7 + 10;
int rt[N], ls[M], rs[M], tot = 0; LL tree[M];
void pushup(int p) { tree[p] = tree[ls[p]] + tree[rs[p]]; }
void update(int &p, int l, int r, int x, int k) {
    if (x > r || x < l) { return ; } if (!p) p = ++ tot;
    if (l == r) { tree[p] += k; return ; } int mid = (l + r) >> 1;
    update(ls[p], l, mid, x, k); update(rs[p], mid + 1, r, x, k); pushup(p); return ;
}
LL query(int p, int l, int r, int x) {
    if (!p || x > r || x < l) { return 0; } if (l == r) return tree[p];
    int mid = (l + r) >> 1; return query(ls[p], l, mid, x) + query(rs[p], mid + 1, r, x); 
}
int merge(int p1, int p2, int l, int r) {
    if (!p1 || !p2) return p1 + p2;
    if (l == r) { tree[p1] += tree[p2]; tree[p2] = 0; return p1; }
    int mid = (l + r) >> 1;
    ls[p1] = merge(ls[p1], ls[p2], l, mid); rs[p1] = merge(rs[p1], rs[p2], mid + 1, r);
    pushup(p1); tree[p2] = ls[p2] = rs[p2] = 0; return p1;
}
}
void DFS2(int u, int f) {
    for (int v : F[u]) {
        DFS2(v, u); Sgt :: rt[u] = Sgt :: merge(Sgt :: rt[u], Sgt :: rt[v], 1, n);
    } if (u <= n) Sgt :: update(Sgt :: rt[u], 1, n, depth[u], W[u]);
    for (auto [v, id] : Qry[u]) 
        Ans[id] = Sgt :: query(Sgt :: rt[u], 1, n, depth[u] + v);
    return ;
}

// Find k-th
int Jump(int x, int k) {
    for (int i = 20; i >= 0; i --) if ((k >> i) & 1) x = fa[i][x];
    return x;
}

int main() {
    ios :: sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n >> Q;
    for (int i = 1; i <= n; has[col[i]] = 1, i ++) cin >> col[i];
    for (int i = 1; i <= n; i ++) cin >> W[i];
    for (int i = 2, u; i <= n; i ++) {
        cin >> u; G[u].emplace_back(i);
    } DFS(1, 0); int x, y;
    for (int k = 1; k <= 20; k ++) for (int i = 1; i <= n; i ++)
        fa[k][i] = fa[k - 1][fa[k - 1][i]];
    for (int i = 1; i <= Q; i ++) {
        cin >> x >> y; int l = Jump(x, y);
        if (!l) Ans[i] = 0;
        else Qry[l].emplace_back(make_pair(y, i));
    }
    for (int i = 1; i <= 500000; i ++) if (has[i])
        DFS2(n + i, 0);
    for (int i = 1; i <= Q; i ++) cout << Ans[i] << "\n";
    return 0;
}