2019牛客国庆集训派对day3 J 买一送一 （思维+树上dfs+组合计数）

大致题意

ICPCCamp 有 n 个商店，用 1,2,…,n 编号。对于任意 i > 1，有从商店 pi ​ 到 i 的单向道路。
同时，商店 i 出售类型为 ai​ 的商品。
Bobo 从商店 1 出发前往商店 i。他要在两个不同的商店购买商品（包括商店 1 和 i）。设他先购买的商品类型是 x，后购买的商品类型是 y，他用 fif_ifi​ 表示不同的有序对 ⟨x,y⟩ 的数量。
求出 f2,f3,…,fn​ 的值。

思路

由于除去 1 以外每个点只有一个点能够到它，所有从 1 好点出发能遍历到的恰好是一棵树。剩余无法遍历到的点可能是一些环伸出一些边，或者也是一些树。但是只要有一个环，图就会不连通，多出几个连通块。因为Bobe只会从 1 号节点出发，所以只要知道 1 号节点为根的是一棵树即可。然后考虑如何统计答案。如果只是一个序列，我们可以从前往后统计，同时记录到当前点 i ，共有多少种不同的种类 sum[i]，用 pre[i]表示前一个种类与 i 相同的位置，那么这一次的答案就是增加 pre[i] 到 i-1 之间新增的种类数，与 a[i] 进行搭配。
所以 ans[i] = ans[i-1] + sum[i-1] - sum[pre[i]-1]；
树形结构就dfs遍历整棵树维护即可，把深度当作序列即可。回溯的时候清空信息。

代码

比赛的时候写的急，代码有点丑陋…

#include<bits/stdc++.h>
using namespace std;
#define maxn 100005
#define maxm 1006
#define ll long long int
#define INF 0x3f3f3f3f
#define inc(i,l,r) for(int i=l;i<=r;i++)
#define dec(i,r,l) for(int i=r;i>=l;i--)
#define mem(a) memset(a,0,sizeof(a))
#define sqr(x) (x*x)
#define inf (ll)2e18+1
#define PI acos(-1)
#define mod 10007
#define auto(i,x) for(int i=head[x];i;i=ed[i].nxt)
ll read(){
    ll x=0,f=1ll;char ch=getchar();
    while(!isdigit(ch)){if(ch=='-')f=-1;ch=getchar();}
     while(isdigit(ch))x=x*10+ch-'0',ch=getchar();
     return f*x;
}
int n,a[maxn],sum[maxn],pre[maxn],vis[maxn],ans[maxn],dep[maxn];
struct edge{int v,nxt;}ed[maxn<<1];
int head[maxn],tot;
void add(int x,int y){ed[++tot]={y,head[x]};head[x]=tot; }
void ed_clr(int n){tot=0;inc(i,1,n)head[i]=0; }
void dfs(int x,int y){
    dep[x]=dep[y]+1;
    ans[x]+=ans[y];
    sum[dep[x]]+=sum[dep[y]];
    int xx=((pre[a[x]]==0)?0:(pre[a[x]]-1));
    ans[x]+=sum[dep[x]]-sum[xx];
    int pp=pre[a[x]];
    if(vis[a[x]]==0){
        sum[dep[x]]++;
    }
    pre[a[x]]=dep[x];
    vis[a[x]]++;
    for(int i=head[x];i;i=ed[i].nxt){
        int v=ed[i].v;
        dfs(v,x);
    }
    sum[dep[x]]=0;
    pre[a[x]]=pp;
    vis[a[x]]--;
}
int main()
{
    while(~scanf("%d",&n)){
        int x;
        inc(i,2,n)add(read(),i);
        inc(i,1,n)a[i]=read();
        dfs(1,0);
        inc(i,2,n)printf("%d\n",ans[i]);
        inc(i,0,n)ans[i]=head[i]=sum[i]=0;
        inc(i,1,n)pre[a[i]]=vis[a[i]]=0;
        tot=0;
    }
    return 0;
}