[ACM]【容斥原理/树形DFS】Atcoder163 path pass i

path pass i

题意：有一棵树，每个点有不同颜色（也有相同）。计算包括每种颜色的点的路径数目。
注意颜色的序号范围是1至N，与点的数目是一样的，所以输出（除非颜色数目与点数目相等）肯定有一些答案为0。

思路：

因为本人太菜，补这道题补到现在才会才A。
首先，最重要的思路是容斥原理。
经过思考后会发现，如果算包含一种颜色的点的路径，是很困难的。（1-2-3包含2，2-3也包含2，2也包含2，1-2也包含……要是点的数目大，似乎无规律可循）
同时，我们知道，算一个连通块的路径的数目，就是n*(n-1)/2（排列组合Cn2），注意到这道题中，只有一个点也算一条路径，那么加上n，得到算一个连通块的路径的数目为n*(n+1)/2。那么问题就可以转化为，求“路径总数-不包含颜色为i的点的所有连通块内部的路径数”。
接下来就只用找出不包含每种颜色的所有连通块的大小。

思路似乎不难，但是实现真是把我虐死了
实现是用了树形DFS，一遍把树全部遍历（O(n)），答案就出来了。
具体操作：
（1）存树：可以用vector存一个点所能到达的所有点；也可以用前向星（我不确定这个是不是）结构体存边。要注意是无向图，存边要存两次！！
（2）DFS纵向：一直DFS到最底层之后，回溯的过程非常重要：从叶子往根走的过程中，每遇到一种颜色，马上计算上一次（也是在回溯过程中）遇到的这种颜色离这里的位置，就可以得出之间并非此颜色的节点数目。更新一次答案。把这种颜色的最近位置更新（怎么存位置，见代码）。那么回溯到根节点的时候，除了根节点颜色以外的其他颜色都没有进行最后一次更新（根附近的并非此颜色的连通块没有被计算），这就在DFS结束之后补上。
（3）DFS横向：因为是树嘛，可能不是一条直链，而有很多分链，DFS的时候就要考虑到这些分链的作用。于是每到DFS一个节点的时候，都要记录当前的数目，每次进入DFS之前，记录一下，回来之后，计算一下改变量。

代码：

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=200005;
int tot,color[maxn],head[maxn],sz[maxn],sum[maxn];
/*
tot：边的总数
color：存颜色，head：存边头
sz：size存树的大小（i的子树（包括i）的大小）
sum：颜色为i的节点的子树大小
*/
ll ans[maxn];
struct node{
	int to,next;
}edge[maxn*2];
void add_edge(int u,int v){
	++tot;
	edge[tot].to=v;
	edge[tot].next=head[u];
	head[u]=tot;
}
ll f(ll n){return n*(n+1)/2;}
void dfs(int u,int fa){	
	int c=color[u],last,delta,v;
	sz[u]=1;
	//这个sum[c]是遍历到u之前的遍历到的颜色为c的节点的子树大小
	ll pre=sum[c];
	for(int i=head[u];i;i=edge[i].next){
		v=edge[i].to;
		if(v==fa) continue;
		//这个for循环内部的上一次dfs得到的sum[c]
		last=sum[c];
		dfs(v,u);
		//这一次dfs遍历中增加的sum[c]量
		delta=sum[c]-last;
		sz[u]+=sz[v];
		//减去上一层颜色为c的节点们离这里的连通块
		ans[c]-=f(sz[v]-delta);
	}
	//更新颜色为c的节点的位置
	//其实也是dfs完了之后，
	//遍历顺序在u之前（包括）得到的c颜色的子树的总的大小
	sum[c]=pre+sz[u];
}
int main(){
	int n;
	scanf("%d",&n);
	for(int i=1;i<=n;i++){
		scanf("%d",&color[i]);
		ans[i]=f(n);
	}
	for(int i=1;i<=n-1;i++){
		int u,v;
		scanf("%d%d",&u,&v);
		add_edge(u,v);
		add_edge(v,u);
	}
	dfs(1,0);
	for(int i=1;i<=n;i++){
		//把离根最近的连通块补上
		ans[i]-=f(n-sum[i]);
		printf("%lld\n",ans[i]);
	}
	return 0;
}