树形DP

题目意思是给一棵树,$n$个点,每个点必须要染$k$个色(即是个颜色的集合),有无限种颜色,但每种颜色最多使用两次,一条边的两个顶点的颜色集合中有相同颜色时,这个边可以算入答案,求最大的答案。

首先,每个颜色最多使用两次,所以相同的颜色放在相邻的两点一定最优,所以我们要选择一些边,使得没有点连接超过$k$条边。$dp[u][f]$ 代表 $u$ 与父亲相连的边是否选择$f=1$代表选择,否则不选择 。$dp[u][0]$ 表示$u$不与父亲相连,我们选取$u$的若干个儿子,由于颜色是无限的,首先 $dp[v][0],v \in son[u]$ 是可以全部加到 $dp[u][0]$ 中去的 , 然后我们可以选择$k$个最大的$dp[v][1] + w$在加到$dp[u][0]$中去,但是选择了$dp[v][1]$就不能选择$dp[v][0]$,所以还要减去$dp[v][0]$,所以我们按照$dp[v][1]+w-dp[v][0]$排个序,取前$k$大即可。$dp[u][1]$ 由于$u$和父亲颜色一样,所以只能取前$k-1$大。

#include <bits/stdc++.h>
using namespace std;
#define lowbit( x ) (x&(-x))
#define debug( x ) cout<<#x<<" is "<<x<<endl;
#define pi acos(-1.0)
#define ls rt<<1
#define rs rt<<1|1
typedef long long ll;
const int N = 5e5 + 10;
const int inf = 0x3f3f3f3f;
const int mod = 6000;
const int M = 2e5;
typedef pair < int , int > P;
typedef pair < ll , int > PL;
vector < P > G[N];
int n , k;
ll dp[N][2];

void dfs ( int u , int fa ) {
    dp[u][0] = dp[u][1] = 0;
    ll sum = 0;
    vector < ll > g;
    for ( auto &E:G[u] ) {
        int v = E.first , w = E.second;
        if ( v == fa ) continue;
        dfs ( v , u );
        sum += dp[v][0];
        g.emplace_back ( dp[v][1] + w - dp[v][0] );
    }
    sort ( g.begin () , g.end () , greater < ll > () );
    int sz = g.size ();
    for ( int i = 0 ; i < min ( k , sz ) ; ++ i ) {
        if ( g[i] > 0 ) sum += g[i];
        else break;
    }
    dp[u][0] = dp[u][1] = sum;
    if ( sz >= k && g[k - 1] > 0 ) dp[u][1] -= g[k - 1];
}

int main () {
    int q;
    scanf ( "%d" , &q );
    while ( q -- ) {
        scanf ( "%d%d" , &n , &k );
        for ( int i = 1 ; i <= n ; ++ i ) G[i].clear ();
        int u , v , w;
        for ( int i = 1 ; i < n ; ++ i ) {
            scanf ( "%d%d%d" , &u , &v , &w );
            G[u].emplace_back ( P ( v , w ) );
            G[v].emplace_back ( P ( u , w ) );
        }
        dfs ( 1 , - 1 );
        printf ( "%lld\n" , dp[1][0] );
    }
    return 0;
}
Last modification:November 8th, 2019 at 08:56 pm
如果觉得我的文章对你有用,请随意赞赏