BZOJ4033 樹上染色

Problem Description

有一棵點數爲N的樹，樹邊有邊權。給你一個在0~N之內的正整數K，你要在這棵樹中選擇K個點，將其染成黑色，並將其他的N-K個點染成白色。將所有點染色後，你會獲得黑點兩兩之間的距離加上白點兩兩之間距離的和的收益。
問收益最大值是多少。

Input

第一行兩個整數N,K。
接下來N-1行每行三個正整數fr,to,dis，表示該樹中存在一條長度爲dis的邊(fr,to)。
輸入保證所有點之間是聯通的。
N<=2000,0<=K<=N

Output

輸出一個正整數，表示收益的最大值。

Example Input

5 2
1 2 3
1 5 1
2 3 1
2 4 2

Example Output

17
【樣例解釋】
將點1,2染黑就能獲得最大收益。

題解

考慮使用樹形DP來做

• DP狀態表示爲dp[x][j],表示在以x爲根的子樹中選取j個點作爲黑點的最大值，siz[x]表示以x爲根的子樹的大小。
• 枚舉子節點y得到一條父親節點與子節點的一條邊，不妨假設邊的權值爲c
• 枚舉子節點上選取黑點的個數爲l,則這條邊產生的貢獻爲兩邊黑點的個數相乘加上兩邊白點的個數相乘再乘上邊的權值，即$c*(l*(k-l)+(siz[y]-l)*(n-k-(siz[y]-l)))$
• 枚舉黑點個數時應從大到小枚舉避免重複計算

C++

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<queue>
#include<set>
#include<map>
#include<stack>

using namespace std;

const long long N=2e3+50;

long long first[N],nxt[N*2],to[N*2],c[N*2],len,n,m,siz[N],dp[N][N],k;

void add(long long x,long long y,long long z)
{
    nxt[++len]=first[x];
    first[x]=len;
    to[len]=y;
    c[len]=z;
}

void dfs(int x,int fa)
{
    siz[x]=1;
    for (int i=first[x];i;i=nxt[i])
    if (to[i]!=fa)
    {
        int y=to[i];
        dfs(y,x);

        for (int j=min(siz[x],k);j>=0;j--)
        {
            for (int l=min(siz[y],k);l>=0;l--)
            {
                dp[x][j+l]=max(dp[x][j+l],dp[y][l]+dp[x][j]+c[i]*(l*(k-l)+(siz[y]-l)*(n-k-(siz[y]-l))));
            }
        }
        siz[x]+=siz[y];
    }
}

signed main()
{
#ifdef ONLINE_JUDGE
#else
    freopen("r.txt", "r", stdin);
#endif

    cin>>n>>k;
    for (int a,b,c,i=1;i<n;i++)
    {
        cin>>a>>b>>c;
        add(a,b,c);
        add(b,a,c);
    }

    dfs(1,0);

    cout<<dp[1][k]<<endl;
}