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;
}