D. Weight the Tree---树形dp

题目链接

#include <iostream>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;
const int N = 200010,mod = 1e9+7;

vector<int>g[N];
int w[N];
struct node
{
	int g,s;//好节点的数量和权值之和 
}f[N][2];

bool operator > (node &a,node &b)
{
	if(a.g!=b.g)return a.g>b.g;
	return a.s<b.s;
} 

void dfs1(int u,int fa)
{
	f[u][1].g=1;
	f[u][1].s=g[u].size();
	f[u][0].s=1;
	for(auto v : g[u])
	{
		if(v==fa)continue;
		dfs1(v,u);
		f[u][1].g+=f[v][0].g;
		f[u][1].s+=f[v][0].s;
		int best=(f[v][0]>f[v][1]?0:1);
		f[u][0].g+=f[v][best].g;
		f[u][0].s+=f[v][best].s;
	}
}

void dfs2(int u,int fa,int b)
{
	w[u]=(b?g[u].size():1);
	for(auto v : g[u])
	{
		if(v==fa)continue;
		if(b)dfs2(v,u,0);
		else
		{
			int best=(f[v][0]>f[v][1]?0:1);
			dfs2(v,u,best);
		}
	}
}

signed main()
{	
	int n;cin>>n;
	for(int i=1,a,b;i<n;i++)
	{
		cin>>a>>b;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	if(n==2){cout<<"2 2"<<'\n';cout<<"1 1"<<'\n';return 0;}
	dfs1(1,-1);
	int best=(f[1][0]>f[1][1]?0:1);
	dfs2(1,0,best);
	cout<<f[1][best].g<<" "<<f[1][best].s<<'\n';
	for(int i=1;i<=n;i++)cout<<w[i]<<" ";
	cout<<'\n';
}