Codeforces Round #591 (Div. 2)E. Paint the Tree(树形dp+贪心(优先队列))

E. Paint the Tree

题意：

输入$n,k(5e5)$；
接下来$n-1$行，每行$u,v,w$,表示树边；
每个点有$k$个颜色，每种颜色出现次数不超过两次，当一条边两端有颜色相同，这条边的权值算作贡献，问贡献和最大为多少？

题解：

$dp[u][0]$表示$u$子树不和父亲节点有相同颜色的贡献和最大值；
$dp[u][1]$表示$u$子树和父亲节点有相同颜色的贡献最大值；
那么如何转移？
对于节点$u$,有$v_1,v_2,\dots ,v_x$儿子
贪心的选择比较优的儿子就好，注意用优先队列维护。

代码：

#include
using namespace std;
const int N=5e5+7;
#define ll long long
int q,n,k;
struct Edge{
    int v,w,nxt;
}e[N<<1];
int head[N],cnt;
inline void add(int u,int v,int w){
    e[cnt]=(Edge){v,w,head[u]};
    head[u]=cnt++;
}
ll dp[N][2];
void dfs(int u,int fa){
    priority_queue<ll>q;ll sum=0;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(e1.v==fa)continue;
        dfs(e1.v,u);
        ll t=dp[e1.v][1]+e1.w-dp[e1.v][0];
        if(t>0)q.push(t);
        sum+=dp[e1.v][0];
    }
    if(k-1==0)dp[u][1]=sum;
    int cnt=0;ll he=0;
    while(!q.empty()){
        he+=q.top();q.pop();cnt++;
        if(cnt==k-1)dp[u][1]=sum+he;
        if(cnt==k)dp[u][0]=sum+he;
    }
    if(cnt<k-1)dp[u][1]=sum+he;
    if(cnt<k)dp[u][0]=sum+he;
}
void solve(){
    for(int i=1;i<=n;i++)dp[i][0]=dp[i][1]=0;
    dfs(1,0);
    printf("%lld\n",max(dp[1][0],dp[1][1]));
}
int main(){
  //  freopen("tt.in","r",stdin),freopen("tt.out","w",stdout);
    scanf("%d",&q);
    while(q--){
        scanf("%d%d",&n,&k);
        for(int i=1;i<=n;i++)head[i]=-1;
      //  memset(head,-1,sizeof(head));就是因为这个超时了。
        cnt=0;
        for(int i=1;i<n;i++){
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w),add(v,u,w);
        }
        solve();
    }
    return 0;
}