E. Paint the Tree(树形dp+贪心)

题目

题意：

    给定一棵带权树，你必须给每个点涂上k种颜色，每种颜色至多用两次。对于一条边连着的两个点，如果有一种颜色相同，那么就会对答案产生边权的贡献。要求分配颜色使得边权最大。
    $1\le n,k\le 5\cdot 10^5,1\le u_i,v_i\le n,u_i\ne v_i,1\le w_i\le 10^5$

分析：

    理解一下题意就是每个点最多有k个机会与相邻的点相连对答案产生贡献。对于一个点来说，它要与另一个点相连，那么另一个点能连的边数就少了，所以不能贪心的选择边权大的边。
    我们发现给定的是一棵树，贪心不行自然就想dp。又是一棵树，就可以往树形dp的方向上想。对于一个点，它与其父亲节点连与不连两种选择。dp[i][0]表示与父亲不连以i为根的子树的最大值。那么有k个可连，连一个产生的多出来的贡献为dp[t][1]+v-dp[t][0]。排序选最大的k个即可。dp[i][1]选k-1个。

#include <iostream>
#include <vector>
#include <algorithm> 
using namespace std;

typedef long long ll;

struct node{
	int id,v;
	node(int a,int b)
	{
		id = a;
		v = b;
	}
};

vector<node> g[500005];

ll dp[500005][2];
int tot;

void dfs(int x,int fa)
{
	vector<ll> add;
	dp[x][1] = dp[x][0] = 0;
	for (int i = 0; i < g[x].size(); i++)
	{
		node t = g[x][i];
		if( t.id == fa ) continue;
		dfs(t.id,x);
		dp[x][0] += dp[t.id][0];
		add.push_back(dp[t.id][1]+t.v-dp[t.id][0]); 
	}
	dp[x][1] = dp[x][0];
	sort(add.begin(),add.end());
	int rest1 = tot;
	int rest2 = tot - 1;
	for (int i = add.size() - 1; i >= 0; i--)
	{
		if( add[i] < 0 ) break;
		if( rest1 )
		{
			dp[x][0] += add[i];
			rest1 --; 
		}
		if( rest2 )
		{
			dp[x][1] += add[i];
			rest2 --;
		}
	}
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	int t;
	cin >> t;
	while( t-- )
	{
		int n,k;
		cin >> n >> k;
		tot = k;
		for (int i = 1; i <= n; i++) g[i].clear();
		for (int i = 1; i < n; i++)
		{
			int x,y,v;
			cin >> x >> y >> v;
			g[x].push_back(node(y,v));
			g[y].push_back(node(x,v));  
		}
		dfs(1,0);
		cout << dp[1][0] << '\n';
	}
	return 0;
}