iai 彩色树 题解

前置题目：树的颜色（时序差分）。

每种颜色的贡献为至少出现一次的路径数。容斥，总路径数减去未出现过的路径数。

记 $f_c(u)$ 表示以 $u$ 为根子树内，结点到 $u$ 路径上无颜色 $c$ 的结点个数。

钦定统计颜色为 $c$ 的贡献，
$$\sum _{co{l}_u=c}\frac{f_c(u)(f_c(u)-1)}{2}+\frac{f_c(1)(f_c(1)-1)}{2}$$

可以证明这样统计是不重不漏的。

$f$ 显然可以通过时序差分降维。

实现可能会有个简单的标记永久化思想，或者通过状态转成补集避免标记。

#include
#define int long long
using namespace std;

const int N = 2e5 + 5;

int n, col[N], p[N], ans, f[N], g[N], tag;

vector <int> son[N];

int s(int x){return x*(x-1)/2;}

void dfs(int u)
{
	f[col[u]] --, tag ++ ;
	for(auto v : son[u])
	{
		int before = f[col[u]] + tag;
		dfs(v);
		int after = f[col[u]] + tag;
		int cur = after - before;
		ans -= s(cur); f[col[u]] -= cur;
	}
}

signed main()
{
	ios::sync_with_stdio(0); cin.tie(0);
	
	cin >> n;
	
	for(int i=1; i<=n; i++) cin >> col[i];
	
	for(int i=2; i<=n; i++) cin >> p[i], son[p[i]].push_back(i);
	
	ans = n * s(n);
	
	dfs(1);
	
	for(int i=1; i<=n; i++) 
		ans -= s(f[i] + tag);
		
	cout << ans;
}