Codeforces 246E Blood Cousins Return

Codeforces 246E Blood Cousins Return

數據結構，dsu on tree

題意

一棵樹，每一個點有一個顏色（字符串），每一次詢問以某一個點爲根的子樹中與其距離爲k的點有多少種顏色。

思路

dsu on tree。因爲第二次dfs時每個深度的信息都要記錄，就對每一個深度開一個set，記錄這個深度下出現過的顏色即可。
dfs完一個節點，將它的深度加上詢問的深度，就是實際深度，查詢這個深度的set的size即可。
注意不是樹，是森林，記一下所有的root即可。

代碼

#include<bits\stdc++.h>
#define M(a,b) memset(a,b,sizeof(a))
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
using namespace std;
const int MAXN=100007;
typedef long long LL;
struct Edge
{
    int to, ne;
}e[MAXN<<1];
int head[MAXN], edgenum;
int val[MAXN];
struct Query
{
    int k, ne;int id;
}qn[MAXN];
int qhead[MAXN], qam;
int sz[MAXN], hson[MAXN], dp[MAXN];
map<string, int> ha;
void findson(int u, int fa)
{
    sz[u]=1;
    if(fa==-1) dp[u]=1;
    else dp[u]=dp[fa]+1;
    for(int i=head[u];~i;i=e[i].ne)
    {
        int to=e[i].to;
        if(to!=fa)
        {
            findson(to, u);
            if(!hson) hson[u]=to;
            else if(sz[hson[u]]<sz[to]) hson[u]=to;
            sz[u]+=sz[to];
        }
    }
}
int hs;set<int> cnt[MAXN];
void suan(int u, int fa, int v)
{
    if(v==1)
        cnt[dp[u]].insert(val[u]);
    else cnt[dp[u]].erase(val[u]);
    for(int i=head[u];~i;i=e[i].ne)
        if(e[i].to!=fa&&e[i].to!=hs)
            suan(e[i].to, u, v);
}
int res[MAXN];
void dfs(int u, int fa, int kp)
{
    for(int i=head[u];~i;i=e[i].ne)
    {
        int to=e[i].to;
        if(to!=fa&&to!=hson[u])
            dfs(to, u, 0);
    }
    if(hson[u]) dfs(hson[u], u, 1), hs=hson[u];
    suan(u, fa, 1);hs=0;
    for(int i=qhead[u];~i;i=qn[i].ne)
    {
        int hi=dp[u]+qn[i].k;
        if(hi>=MAXN) res[qn[i].id]=0;
        else res[qn[i].id]=cnt[hi].size();
    }
    if(!kp) suan(u, fa, -1);
}
vector<int> roots;
int main()
{
    int n;
    while(scanf("%d", &n)==1)
    {
        M(head, -1), M(qhead, -1);edgenum=qam=0;
        roots.clear();
        int colam=0;
        for(int i=1;i<=n;i++)
        {
            string tmps;cin>>tmps;
            if(!ha[tmps]) ha[tmps]=++colam, val[i]=colam;
            else val[i]=ha[tmps];
            int p=i, q;scanf("%d", &q);
            if(q==0)
            {
                roots.push_back(i);
                continue;
            }
            e[edgenum].to=q, e[edgenum].ne=head[p], head[p]=edgenum++;
            e[edgenum].to=p, e[edgenum].ne=head[q], head[q]=edgenum++;
        }
        int q;scanf("%d", &q);
        for(int i=1;i<=q;i++)
        {
            int t1, t2;scanf("%d%d", &t1, &t2);
            qn[qam].k=t2, qn[qam].ne=qhead[t1], qn[qam].id=i, qhead[t1]=qam++;
        }
        for(int i : roots)
            findson(i, -1);
        for(int i:roots)
            dfs(i, -1, 0);
        for(int i=1;i<=q;i++)
            printf("%d\n", res[i]);
    }
    //system("pause");
    return 0;
}