codeforces1467E Distinctive Roots in a Tree

题目链接

题目：

给定一棵$n$个结点的树，每个点有一个权值$a_i$，问有多少个点$u$满足从$u$出发到其他所有点的路径中都不存在两个权值相同的点。
$(1\le n\le 2\times 10^5,1\le a_i\le 10^9)$

题解：

一开始拿到这个题我想的是用换根$dp$，但发现有一些信息维护不了。
这道题应该从性质入手，先随便确定一个根，考虑有两个权值相同的点的情况，设为$u,v$：
（1）$u,v$存在祖先关系
   不妨设$u$是$v$的祖先，那么显然可能满足题意的点只能是在$v$所在的那棵以$u$的某个子节点为根的子树中的而又不在以$v$为根的子树中的点。觉得不清楚可以看图：

（2）$u,v$不存在祖先关系
   那么可以确定在以$u$为根的子树中和以$v$为根节点的子树中的点是不满足题意的。
那么上述的判法将所有不满足题意的点都判掉了吗？答案是肯定的，因为如果一个点不满足题意，那么这个点到其他所有结点的路径中一定会出现至少两者之一的情况，所以所有不合法的点都会被判掉。
正确性得到证明以后，来考虑具体的算法实现。我们可以用打01标记的方式来维护点的合法性，由于有多次更新操作和单次查询操作，所以可以用树上差分来维护。我们用$f_u=0/1$来标识点的合法状态，对于情况（1），我们在遍历$u$的所有子节点，如果在以$v$为根的子树中存在和$a_u$权值相同的点，那么令$f_{root}++,f_v--$，这里具体来讲可以维护一个全局的$map$变量$num$，存下当前$dfs$到的点的权值出现情况，在$dfs$进入$v$子树前存下当前的$nu{m}_{a_u}$，从$v$子树出来的时候比较当前的$nu{m}_{a_u}$和之前存下的是否相等，如果不相等，说明以$a_u$为权值的点在$v$子树中出现了；对于情况（2），也是用一样的思路去维护，如果在以$u$为根节点的子树中没有出现的以$a_u$为权值的点，那么就出现了情况（2），这时将$f_u++$即可。最后$dfs$一遍下顺标记。$f$值为0的点就是满足题意的点。

复杂度：$O(nlogn)$

代码：

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<vector>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<string>
#include<bitset>
#include<sstream>
#include<ctime>
//#include<chrono>
//#include<random>
//#include<unordered_map>
using namespace std;

#define ll long long
#define ls o<<1
#define rs o<<1|1
#define pii pair<int,int>
#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define sz(x) (int)(x).size()
#define all(x) (x).begin(),(x).end()
const double pi=acos(-1.0);
const double eps=1e-6;
const int mod=1e9+7;
const int INF=0x3f3f3f3f;
const int maxn=2e5+5;
ll read(){
	ll x=0,f=1;
	char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}
int n,ans;
vector<int>g[maxn];
int f[maxn],a[maxn];
map<int,int>num,sta;
void dfs(int u,int fa){
	int base=sta[a[u]];
	sta[a[u]]++;
	for(auto v:g[u]){
		if(v==fa)continue;
		int tmp=sta[a[u]];
		dfs(v,u);
		if(tmp!=sta[a[u]]){
			f[1]++;
			f[v]--;
		}
	}
	int diff=sta[a[u]]-base;
	if(diff!=num[a[u]]){
		f[u]++;
	}
}
void solve(int u,int fa,int x){
	if(!x){
		++ans;
	}
	for(auto v:g[u]){
		if(v==fa)continue;
		solve(v,u,x+f[v]);
	}
}
int main(void){
	// freopen("in.txt","r",stdin);
	scanf("%d",&n);
	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
		num[a[i]]++;
	}
	int u,v;
	for(int i=1;i<=n-1;i++){
		scanf("%d%d",&u,&v);
		g[u].pb(v);
		g[v].pb(u);
	}
	dfs(1,0);
	ans=0;
	solve(1,0,f[1]);
	printf("%d\n",ans);
	return 0;
}