多比特杯武汉工程大学第四届ACM程序设计竞赛同步赛 F能源收集者 （树形dp）

link

直接暴力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;
}