ARC086E - Smuggling Marbles

题目大意

https://www.luogu.com.cn/problem/AT3721

题解

设f[i][j][0/1/2]表示到点i统计的是深度为j的，在点i的和为0/1/>1

长链剖分，再记个时间维护f[2]->f[0]的

code

#include <bits/stdc++.h>
#define fo(a,b,c) for (a=b; a<=c; a++)
#define fd(a,b,c) for (a=b; a>=c; a--)
#define add(a,b) a=((a)+(b))%1000000007
#define min(a,b) (a<b?a:b)
#define max(a,b) (a>b?a:b)
#define mod 1000000007
#define ll long long
//#define file
using namespace std;

int a[200001][2],ls[200001],nx[200001],d[200001],g[200001],sum[200001],len,n,i,j,k,l;
ll p[200001],f[200001][3],F[3],ans;

void New(int x,int y) {++len;a[len][0]=y;a[len][1]=ls[x];ls[x]=len;}
void dfs(int t,int D)
{
	int i;
	++sum[D];
	
	for (i=ls[t]; i; i=a[i][1])
	{
		dfs(a[i][0],D+1);
		if (d[a[i][0]]+1>d[t])
		d[t]=d[a[i][0]]+1,nx[t]=a[i][0];
	}
}
void Dfs(int t)
{
	int i,j,k,l;
	
	f[t][0]=f[t][1]=1;
	if (nx[t]!=-1)
	{
		for (i=ls[t]; i; i=a[i][1]) Dfs(a[i][0]);
		for (i=ls[t]; i; i=a[i][1])
		if (a[i][0]!=nx[t])
		{
			j=nx[t],k=a[i][0];
			while (k!=-1)
			{
				add(f[k][0],f[k][2]),f[k][2]=0;
				if (g[j]!=t)
				g[j]=t,add(f[j][0],f[j][2]),f[j][2]=0;
				
				memset(F,0,sizeof(F));
				add(F[0],f[j][0]*f[k][0]);
				add(F[1],f[j][0]*f[k][1]+f[j][1]*f[k][0]);
				add(F[2],f[j][0]*f[k][2]+f[j][1]*(f[k][1]+f[k][2])+f[j][2]*(f[k][0]+f[k][1]+f[k][2]));
				memcpy(f[j],F,sizeof(F));
				j=nx[j],k=nx[k];
			}
		}
	}
}

int main()
{
	#ifdef file
	freopen("arc086e.in","r",stdin);
	#endif
	
	memset(g,255,sizeof(g));
	memset(nx,255,sizeof(nx));
	scanf("%d",&n);
	p[0]=1;
	fo(i,1,n) scanf("%d",&j),New(j,i),p[i]=p[i-1]*2%mod;
	
	dfs(0,0);
	Dfs(0);
	
	for (i=j=0; i!=-1; add(ans,f[i][1]*p[n+1-sum[j]]),i=nx[i],++j);
	printf("%lld\n",ans);
	
	fclose(stdin);
	fclose(stdout);
	return 0;
}