Codeforces 246E - Blood Cousins Return （树上启发式合并）

246E - Blood Cousins Return

题意

给出一棵家谱树，定义从 u 点向上走 k 步到达的节点为 u 的 k-ancestor，每个节点有名字，名字不唯一。多次查询，给出 u k，问以 u 为根节点的子树下有多少个深度为 dep[u] + k 的节点（dep[u] 为节点 u 的深度）。

分析

和208E - Blood Cousins几乎相同，把那题的 \(C\) 数组换成一个 \(set\) 数组，表示某子树下同一深度所有名字的集合。

code

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int MAXN = 1e5 + 10;
int n;
int fa[MAXN], son[MAXN], dep[MAXN], siz[MAXN];
int col[MAXN];
int cnt, head[MAXN];
struct Edge {
    int to, next;
} e[MAXN << 1];
struct Ex {
    int x, c;
};
vector<Ex> ex[MAXN];
void addedge(int u, int v) {
    e[cnt].to = v; e[cnt].next = head[u]; head[u] = cnt++;
    e[cnt].to = u; e[cnt].next = head[v]; head[v] = cnt++;
}
void dfs(int u) {
    siz[u] = 1;
    son[u] = 0;
    for(int i = head[u]; ~i; i = e[i].next) {
        if(e[i].to != fa[u]) {
            fa[e[i].to] = u;
            dep[e[i].to] = dep[u] + 1;
            dfs(e[i].to);
            if(siz[e[i].to] > siz[son[u]]) son[u] = e[i].to;
            siz[u] += siz[e[i].to];
        }
    }
}
int vis[MAXN], ans[MAXN];
char name[MAXN][22];
set<string> S[MAXN << 1];
void change(int u, int c) {
    if(c > 0) {
        S[dep[u]].insert(name[u]);
    } else S[dep[u]].clear();
    for(int i = head[u]; ~i; i = e[i].next) {
        if(e[i].to != fa[u] && !vis[e[i].to]) change(e[i].to, c);
    }
}
void dfs1(int u, int flg) {
    for(int i = head[u]; ~i; i = e[i].next) {
        if(e[i].to != fa[u] && e[i].to != son[u]) dfs1(e[i].to, 1);
    }
    if(son[u]) {
        dfs1(son[u], 0);
        vis[son[u]] = 1;
    }
    change(u, 1);
    int sz = ex[u].size();
    for(int i = 0; i < sz; i++) {
        ans[ex[u][i].x] = S[ex[u][i].c].size();
    }
    if(son[u]) vis[son[u]] = 0;
    if(flg) change(u, -1);
}
int main() {
    scanf("%d", &n);
    memset(head, -1, sizeof head);
    cnt = 0;
    for(int i = 1; i <= n; i++) {
        int x;
        scanf("%s%d", name[i], &x);
        addedge(i, x);
    }
    dfs(0);
    int m;
    scanf("%d", &m);
    for(int i = 0; i < m; i++) {
        int x, y;
        scanf("%d%d", &x, &y);
        ex[x].push_back(Ex{i, dep[x] + y});
    }
    dfs1(0, 0);
    for(int i = 0; i < m; i++) {
        printf("%d\n", ans[i]);
    }
    return 0;
}