直接暴力set按题意模拟 转化为树的直径
有一点要注意
如果不加
if(a[v].size()>a[u].size())swap(a[u],a[v]);
会mle
#include <bits/stdc++.h>
using namespace std;
#define int long long
//typedef long long ll;
typedef pair<int,int> pii;
#define x first
#define y second
#define pb push_back
#define inf 1e18
#define IOS std::ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);
#define fer(i,a,b) for(int i=a;i<=b;i++)
#define der(i,a,b) for(int i=a;i>=b;i--)
const int maxn=1e5+10;
const int mod=1e9+7;
/*ll qsm(int a,int b)
{ll res=1; while(b){ if(b&1) res=res*a%mod; a=a*a%mod; b>>=1; } return res;}
*/
const int N=1e6+10;
int dr[4][2]={{-1,0},{1,0},{0,-1},{0,1}};
int n,k;
vector<int>g[N];
set<int>a[N];
int ans;
int cnt[N];
int dp[N];
void dfs1(int u,int fa)
{
for(auto v:g[u])
{
if(v==fa)continue;
dfs1(v,u);
if(a[v].size()>a[u].size())swap(a[u],a[v]);
for(auto j:a[v])
{
a[u].insert(j);
}
}
cnt[u]=a[u].size();
}
void dfs2(int u,int fa)
{
dp[u]=cnt[u];
for(auto v:g[u])
{
if(v==fa)continue;
dfs2(v,u);
ans=max(ans,dp[u]+dp[v]);
dp[u]=max(dp[u],dp[v]+cnt[u]);
}
}
void solve()
{
cin>>n>>k;
fer(i,1,n)
{
int x;
cin>>x;
a[i].insert(x);
}
fer(i,1,n-1)
{
int a,b;
cin>>a>>b;
g[a].pb(b);
g[b].pb(a);
}
dfs1(k,-1);
dfs2(k,-1);
cout<<ans<<endl;
}
signed main()
{
IOS;
int _=1;
//cin>>_;
while(_--) solve();
return 0;
}