Codeforces 1241 E Paint the Tree题解（树形DP+贪心）

题目：CF1241E.
题目大意：给定一棵$n$个点的树，每个点可以染$k$种不同的颜色，且每种颜色最多能染给两个点.一条边有贡献当且仅当它的两个端点有至少一种颜色相同，现在问颜色数量无限时的最大贡献和.
数据组数$T\le 5\ast 10^5$，$1\le \sum n,\sum k\le 5\ast 10^5$.

首先一个贪心的思路是把所有边按照边权从大到小排序后依次能选就选，然而这个做法会被这组数据卡掉：

如果$k=1$，这个贪心就假了.

现在我们考虑一个DP的做法，设$f[i][0/1]$表示以$i$为根的子树中，点$i$连向父亲的边选$/$不选的最大贡献和.

那么如何转移呢？我们先假设$f[i][0/1]$全部由$f[so{n}_i][0]$转移过来，之后我们对于$f[i][0]$，我们可以选择其中$k$个儿子把它们变成$f[so{n}_i][1]$，$f[i][1]$同理.

那么考虑求出$f[so{n}_i][1]-f[so{n}_i][0]$的差并排序，贪心选择来转移即可.

时间复杂度$O(\sum n\log n)$.

代码如下：

#include<bits/stdc++.h>
using namespace std;

#define Abigail inline void
typedef long long LL;

const int N=500000;

int n,m;
struct side{
  int y,next,v;
}e[N*2+9];
int lin[N+9],cs;

void Ins(int x,int y,int v){e[++cs].y=y;e[cs].v=v;e[cs].next=lin[x];lin[x]=cs;}
void Ins2(int x,int y,int v){Ins(x,y,v);Ins(y,x,v);}

LL dp[N+9][2],v[N+9],cv;

bool cmp(const LL &a,const LL &b){return a>b;}

void Dfs_dp(int k,int fa){
  dp[k][0]=0;
  for (int i=lin[k];i;i=e[i].next)
    if (e[i].y^fa){
      dp[e[i].y][1]=e[i].v;
      Dfs_dp(e[i].y,k);
      dp[k][0]+=dp[e[i].y][0];
      dp[k][1]+=dp[e[i].y][0]; 
    }
  cv=0;
  for (int i=lin[k];i;i=e[i].next)
    if (e[i].y^fa) v[++cv]=dp[e[i].y][1]-dp[e[i].y][0];
  sort(v+1,v+cv+1,cmp);
  for (int i=1;i<=cv&&i<m&&v[i]>0;++i) dp[k][0]+=v[i],dp[k][1]+=v[i];
  if (m<=cv&&v[m]>0) dp[k][0]+=v[m];
}

Abigail into(){
  scanf("%d%d",&n,&m);
  for (int i=1;i<=n;++i) lin[i]=0;
  for (int i=1;i<n;++i){
  	int x,y,v;
  	scanf("%d%d%d",&x,&y,&v);
  	Ins2(x,y,v);
  }
}

Abigail work(){
  Dfs_dp(1,0);
}

Abigail outo(){
  printf("%I64d\n",dp[1][0]);
}

int main(){
  int T;
  scanf("%d",&T);
  for (;T--;){
    into();
    work();
    outo();
  }
  return 0;
}