iai 彩色树 题解

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

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

考虑颜色 $c$ 的贡献，此时将所有颜色为 $c$ 的点移去，会形成若干连通块。对于大小为 $x$ 的连通块，其贡献为 $\frac{x(x-1)}{2}$。

在这些颜色为 $c$ 的点（以及根节点）处计算对应块的贡献即可。

记 $f_c(u)$ 表示以 $u$ 为根子树内，无颜色 $c$ 的极大连通块大小。

颜色 $c$ 贡献：

$$\sum _{u,co{l}_u=c}\sum _{v\in so{n}_u}\frac{f_c(v)(f_c(v)-1)}{2}+\frac{f_c(1)(f_c(1)-1)}{2}$$

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

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

#include<bits/stdc++.h>
#define int long long
using namespace std;

const int N = 2e5 + 5;

int n, col[N], p[N], ans, f[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;
}