AtCoder abc163_f - path pass i

爬爬爬

这其实算是ABC的F们中比较难的一个

洛谷题目页面传送门 & AtCoder题目页面传送门

有一棵树\(T=(V,E),|V|=n,|E|=n-1\)，节点\(i\)的颜色为\(a_i\in[1,n]\)。求\(\forall i\in[1,n]\)，经过至少\(1\)个颜色为\(i\)的节点的路径个数。

\(n\in\left[1,2\times10^5\right]\)。

先考虑对于某个\(i\)，该如何求答案。注意到，若正面求答案，那么需要容斥一大堆，或者点分治（我还不会），比较麻烦。考虑反面求答案，求不经过任何颜色为\(i\)的节点的路径个数，最后用总路径数\(\dfrac{n(n+1)}2\)减去它即可。考虑“不经过任何颜色为\(i\)的节点的路径个数”怎么算，显然，若将所有颜色为\(i\)的节点删除，会分离成若干CC，设它们的大小分别为\(s_1,s_2,\cdots,s_m\)，那么答案就是\(\sum\limits_{j=1}^m\dfrac{s_j(s_j+1)}2\)，因为一条路径不经过任何颜色为\(i\)的节点当且仅当它的两端点属于同一个CC。

如果只对于一个\(i\)求答案，那么显然可以直接删除所有该删的点，然后随便DFS一下就\(\mathrm O(n)\)出答案了。但现在需要在较短时间内求出对于所有\(i\)的答案，就有那么一丝丝难度了。不难发现，\(\forall i\)，任意一个CC的深度最小的点要么没有父亲（等于根，令\(1\)为根），要么父亲颜色为\(i\)（反证法，如果不成立，那么显然CC可以再往上扩展）。于是我们可以用这个唯一的（唯一性可以反证，如果不唯一，也就是\(2\)个同辈人没有长辈连接，一定不连通）深度最小的点作为整个CC的代表。

又不难发现，任意一个非\(1\)节点，都只能作为一个\(i\)剖出来的CC的代表，因为它们都有一个固定的父亲，也就有一个固定的父亲的颜色，而节点\(1\)却可以作为任意一个\(i\)剖出来的CC的代表，需要特殊开数组处理。所以总CC数为\(\mathrm O(n)\)。考虑一遍DFS求出所有CC的大小。考虑如何算一个CC的大小。先假设整个以\(x\)为根的子树全属于此CC，然后看有多少是被删了的或属于其他CC的，减去即可。显然，在子树内的所有CC代表要么一个为祖先一个为晚辈，此时以它们为根的子树一个包含另一个，要么无关，此时以它们为根的子树不相交，CC代表们的父亲（颜色为\(i\)）亦然。又显然，所有子树内的以颜色为\(i\)的节点为根的子树内的节点都不属于此CC，于是我们只需要将此CC的大小减去所有不为其他颜色为\(i\)的节点的晚辈的颜色为\(i\)的节点的子树大小即可。考虑从下往上贡献，即DFS过程中时刻维护递归栈，每到一个节点就将最深的一个颜色相同的节点（这个可以时刻维护每种颜色的递归栈\(\mathrm O(1)\)做到）在栈中的下一个节点的CC大小减去它的子树大小（正确性显然）。对于根代表多CC的情况，特判一下即可。

时间复杂度\(\mathrm O(n)\)。

下面是AC代码：

#include<bits/stdc++.h>
using namespace std;
#define pb push_back
#define ppb pop_back
typedef long long ll;
const int N=200000;
int n;//节点个数 
int a[N+1];//颜色 
vector<int> nei[N+1];//邻接表 
int sz[N+1]/*子树大小*/,cc_sz[N+1]/*CC大小（非根）*/,cc_sz1[N+1]/*CC大小（根，下标为颜色）*/;
int fa[N+1];//父亲 
vector<int> stk/*总递归栈*/,stk_c[N+1]/*颜色们的递归栈*/;
void dfs(int x=1){//求CC们的大小 
//	cout<<x<<"\n";
	sz[x]=1;
	stk_c[a[x]].pb(stk.size());
	stk.pb(x);//压栈 
	for(int i=0;i<nei[x].size();i++){
		int y=nei[x][i];
		if(y==fa[x])continue;
		fa[y]=x;
		dfs(y);
		sz[x]+=sz[y];
	}
	stk.ppb();
	stk_c[a[x]].ppb();//准备回溯，弹出 
	if(stk_c[a[x]].size()>1)cc_sz[stk[stk_c[a[x]].back()+1]]-=sz[x];//贡献（非根） 
	else/*特判*/ cc_sz1[a[x]]-=sz[x];//贡献（根） 
}
ll ans[N+1];
int main(){
	cin>>n;
	for(int i=1;i<=n;i++)cin>>a[i];
	for(int i=1;i<n;i++){
		int x,y;
		cin>>x>>y;
		nei[x].pb(y);nei[y].pb(x);
	}
	stk.pb(0);
	for(int i=1;i<=n;i++)stk_c[i].pb(0);
	dfs();
	for(int i=1;i<=n;i++)cc_sz[i]+=sz[i],cc_sz1[i]+=n;
//	for(int i=1;i<=n;i++)cout<<cc_sz[i]<<" ";puts("");
//	for(int i=1;i<=n;i++)cout<<cc_sz1[i]<<" ";puts("");
	for(int i=1;i<=n;i++)ans[i]=1ll*n*(n+1)/2;
	for(int i=2;i<=n;i++)ans[a[fa[i]]]-=1ll*cc_sz[i]*(cc_sz[i]+1)/2;
	for(int i=1;i<=n;i++)ans[i]-=1ll*cc_sz1[i]*(cc_sz1[i]+1)/2;//减去反面答案 
	for(int i=1;i<=n;i++)cout<<ans[i]<<"\n";
	return 0;
}