BZOJ4033 [HAOI2015]樹上染色 [樹形DP]

BZOJ4033 [HAOI2015]樹上染色 [樹形DP]

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

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

題解

樹形DP：設f[i][j]$f[i][j]$ 表示以i$i$ 爲根的子樹中有j$j$ 個黑點。有坑點，詳見代碼註釋

一條鏈的收益和，容易想到的是暴力，就是把每條鏈找出來，然後加入Ans。但這道題光點就有2000，所以要考慮更優秀的算法，即不能從找鏈入手。

可以考慮從線段被使用了多少次，即線段對答案的總貢獻入手，這樣就不需要找出鏈具體是哪些。而現在已知以i$i$ 爲根的子樹內有j$j$ 個黑點，那麼這棵子樹以外的黑點數也是已知的，爲k−j$k-j$ 。那麼對於邊e$e$ ，它對答案的貢獻是：

Len[e]∗內部黑點數∗外部黑點數+Len[e]∗內部白點數∗外部白點數$$Len[e]\ast 內部黑點數\ast 外部黑點數+Len[e]\ast 內部白點數\ast 外部白點數$$

狀態……看代碼吧……

經驗

• 分解步驟。一條鏈對答案的貢獻可以看做每一個線段對答案的貢獻；一個數對答案的貢獻可以看做每一個二進制數數位對答案的貢獻。
• 減少枚舉上限。防止時間複雜度退化。
• 子樹兩兩合併。

代碼

#include<cstdio>
#include<iostream>
#define NN 2100
#define ll long long
using namespace std;
ll N,K;
ll Size[NN],f[NN][NN];
ll End[NN<<1],Last[NN<<1],Next[NN<<1],Len[NN<<1],cnt;
void DFS(ll u,ll fa){
    Size[u]=1;
    for(ll i=Last[u];i;i=Next[i]){
        ll v=End[i];
        if(v==fa)continue;
        DFS(v,u);
        //注意此處不要從n開始枚舉，否則時間複雜度會充O(n^2)退化成O(n^3)
        for(ll j=min(Size[u],K);j>=0;j--)
            //j表示的是以u爲根的子樹中有多少黑點（注意：不包含以v爲根的子樹！）
            //而且目前的統計還不包含沒討論過的子樹，即v和v以前的子樹（這一點從Size的更新就可以看出）

            //一定要倒着枚舉，因爲這樣纔不會在這一輪更新中，讓先更新的答案影響到後更新的答案
            //如果非要正着枚舉，也可以開一個臨時數組
            //總的來說，就是要防止重複更新
            //這是因爲f數組的特性：它在更新完以v爲根的這棵子樹前，答案都是v以後的子樹貢獻的
            //因此要把以v爲根的子樹更新完之後才能最終更新f數組的答案（或者在不影響其它答案的情況下更新f數組）
            for(ll k=min(Size[v],K);k>=0;k--){
                //k表示的是以v爲根的子樹中有多少黑點
                if(j+k>K)continue;
                ll t=Len[i]*(k*(K-k)+(Size[v]-k)*(N-K-Size[v]+k));
                //一條邊對答案的貢獻（詳見“題解”）
                f[u][j+k]=max(f[u][j+k],f[u][j]+f[v][k]+t);
            }
        Size[u]+=Size[v];
    }
}
void Ins(ll x,ll y,ll w){
    End[++cnt]=y,Len[cnt]=w;
    Next[cnt]=Last[x],Last[x]=cnt;
}
int main(){
    scanf("%lld%lld",&N,&K);
    for(ll i=1;i<N;i++){
        ll u,v,w;scanf("%lld%lld%lld",&u,&v,&w);
        Ins(u,v,w),Ins(v,u,w);
    }
    DFS(1,0);
    printf("%lld",f[1][K]);
    return 0;
}