2019牛客国庆集训派对day3 J.买一送一(dfs+组合数学)

Posted 2021-05-18 issue是fw

LINK

从根节点$1$开始$dfs$

我们保存一个$sum$表示路径上出现了$sum$种不同的颜色

当我们$dfs$到$u$点时,显然前面$sum$种颜色都可以和$p[u]$形成一个合法的对

然而会有重复,也就是我们之前可能已经算过$p[u]$颜色作为对的情况!!

那么就先减去之前的,再加上现在的

然后当$dfs$退出点$u$时,需要把所有数组,变量都还原回去

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int maxn = 2e5+10;
vector<int>vec[maxn];
int f[maxn],col[maxn],ans,sum,p[maxn],val[maxn],n;
void dfs(int u,int fa)
{
	int temp = val[p[u]] , now = sum;//1到u路径之前最晚出现p[u]的地方 
	ans -= temp; val[p[u]] = now; ans += now;
	f[u] = ans;
	if( ++col[p[u]]==1 )	sum++;
	for( auto v:vec[u] )
		if( v!=fa )	dfs(v,u);
	if( --col[p[u]]==0 )	sum--;
	ans -= now, val[p[u]] = temp, ans += temp;
}
signed main()
{
	while( cin >> n )
	{
		for(int i=2;i<=n;i++)	
		{
			int x; cin >> x;
			vec[x].push_back( i );
		}
		for(int i=1;i<=n;i++)	cin >> p[i];
		dfs(1,1);
		for(int i=2;i<=n;i++)	cout << f[i] << endl;

		ans = sum = 0;
		for(int i=1;i<=n;i++)	vec[i].clear(),f[i] = val[i] = col[i] = 0;
	}
}