「算法筆記」樹形 DP

一、樹形 DP 基礎

又是一篇鴿了好久的文章……以下面這道題為例，介紹一下樹形 DP 的一般過程。

POJ 2342 Anniversary party

題目大意：有一家公司要舉行一個聚會，一共有 \(n\) 個員工，其中上下級的關系通過樹形給出。每個人都不想與自己的直接上級同時參加聚會。每個員工都有一個歡樂度，舉辦聚會的你需要確定邀請的員工集合，使得它們的歡樂度之和最大，並且沒有一個受邀的員工需要與他的直接上級共同參加聚會。\(n\leq 6000\)。

Solution：

考慮一個子樹往上轉移，發現除了子樹的根選與不選的狀態對上面的決策有影響之外，子樹中其他的節點的狀態都不用考慮。

設 \({dp}_{i,j}\) 表示以 \(i\) 號節點為根的子樹，\(j\) 表示第 \(i\) 號節點選或不選的狀態（比如 \(0\) 表示不選，\(1\) 表示選）時，最大的子樹中受邀的人的歡樂度之和。

\({dp}_{u,0}=\sum\limits_{v\in son(u)} \max({dp}_{v,0},{dp}_{v,1})\)（上級不參加舞會時，下級可以參加，也可以不參加）

\({dp}_{u,1}=a_u+\sum\limits_{v\in son(u)} {dp}_{v,0}\)（上級參加舞會時，下級都不會參加）

最后的答案就是 \(\max({dp}_{root,0},{dp}_{root,1})\)，時間復雜度 \(O(n)\)。

void dfs(int x,int fa){
    f[x][0]=0,f[x][1]=a[x];    //這里的 f 數組就是之前講的 dp 數組 
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x),f[x][0]+=max(f[y][0],f[y][1]),f[x][1]+=f[y][0];
    }
}

普通的樹形 dp 中，常常會采用葉→根的轉移形式，若子節點有多個，則需要一一枚舉，將子節點（子樹）的 dp 值合並。dp 的狀態表示中，第一維通常是節點編號（代表以該節點為根的子樹）。大多數時候，我們采用遞歸的方式實現樹形 dp。

二、處理樹上問題的基礎

1. 樹的重心

定義：樹的重心也叫樹的質心。對於一棵 \(n\) 個節點的無根樹，找到一個點，使得把樹變成以該點為根的有根樹時，最大子樹的節點數最小。換句話說，刪除這個點后最大連通塊（一定是樹）的節點數最小。

性質：

1. 一棵樹最多有兩個重心，如果有兩個重心，它們必定有一條邊相連。

2. 樹中所有點到某個點的距離和中，到重心的距離和是最小的，如果有兩個重心，它們的距離和一樣。

3. 把兩棵樹通過一條邊相連，新的樹的重心在原本兩棵樹重心的連線上。

4. 一棵樹添加或者刪除一個節點，樹的重心最多只移動一條邊的位置。

求法：把樹上的節點 \(u\) 刪除后，連通塊為所有 \(u\) 的每個兒子的子樹以及 \(u\) 的父親連出去的整個連通塊。如下圖所示：

考慮 DFS 計算出每棵子樹的節點個數。記 \(sz_x\) 為以 \(x\) 為根的子樹大小。對於 \(\forall v\in son(u)\)，我們都可以通過 \(sz_v\) 知道它的子樹大小。那么 \(u\) 的父親連出去的整個連通塊的大小就是 \(n-sz_u\)，其中 \(n\) 為總節點數。所以我們可以直接計算刪除每個點后的最大連通塊的節點數。

於是我們可以枚舉每個點，找到刪除這個點后最大連通塊的節點數最小的節點。

代碼片段：

void dfs(int x,int fa){
    sz[x]=1,mx[x]=0;    //sz[x]:以 x 為根的子樹大小。  mx[x]:刪除點 x 后的最大連通塊的節點數。 
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x),sz[x]+=sz[y],mx[x]=max(mx[x],sz[y]);
    }
    mx[x]=max(mx[x],n-sz[x]);
    if(mx[x]<ans) ans=mx[x],id=x;    //找到刪除這個點后最大連通塊的節點數最小的節點。id:重心編號。注意 ans 初始化為無窮大。 
}

2. 樹的直徑

定義：樹中兩點間的最長路徑。（樹的直徑可能有很多條）

有一些不同的求法。

（1）通過 LCA 找出樹的一條直徑。

顯然，一條直徑上的所有點有一個共同的 \(\text{LCA}\)。在 DFS 的過程中對於每一個點，考慮以它為 \(\text{LCA}\) 的可能的路徑。

維護以每個點為頂端的最長鏈和次長鏈，然后用最長鏈加上次長鏈更新直徑即可。

相關代碼如下：

int dfs(int x,int fa){
    int mx=0,mx2=0;    //mx:最長鏈長度。mx2:次長鏈長度。 
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        int k=dfs(y,x);
        if(k>mx) mx2=mx,mx=k;
        else if(k>mx2) mx2=k;
    }
    ans=max(ans,mx+mx2);    //最長鏈加次長鏈 
    return mx;
}

（2）通過兩次遍歷找出樹的一條直徑。

第一次遍歷，找出距離某個節點（例如根節點）最遠的一個點 \(u\)。

第二次遍歷，找出距離節點 \(u\) 最遠的一個點 \(v\)。

\(u\) 到 \(v\) 的簡單路徑，即為樹的一條直徑。

另外，為了找出距離某個點最遠的點，這棵樹應該看作無根樹，一個節點連向父親的邊也要存入鄰接表中。

相關代碼如下：（這種方法適用於邊權非負的情況）

void dfs(int x,int fa){
    dep[x]=dep[fa]+1;    //計算每個點的深度 
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y!=fa) dfs(y,x);
    } 
} 
void solve(){
    dfs(1,0),x=1;
    for(int i=2;i<=n;i++)
        if(dep[i]>dep[x]) x=i;    //找出距離根節點最遠的一個點 x 
    dfs(x,0),y=1;
    for(int i=2;i<=n;i++)
        if(dep[i]>dep[y]) y=i;    //找出距離節點 x 最遠的一個點 y 
    printf("%lld %lld\n",x,y);    //x 到 y 的簡單路徑，即為樹的一條直徑 
    printf("%lld\n",dep[y]);    //dep[y] 即樹的直徑的長度 
}

（3）樹形 dp 求樹的直徑

令 \(f_i\) 表示以 \(i\) 為根，到它子樹的葉節點的最大距離。

\(f_u=\max\limits_{v\in son(u)}\{f_v+dis(u,v)\}\)

\(Ans=\max\{f_u+f_v+dis(u,v)\}\)

另外，因為要用當前的 \(f_u\) 更新答案，所以要先更新 \(Ans\) 再更新 \(f_u\)。

void dfs(int x,int fa){
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x),ans=max(ans,f[x]+f[y]+val[i]),f[x]=max(f[x],f[y]+val[i]);    //轉移。其中 val[i] 表示邊 i 的邊權。 
    }
}

三、樹形背包

Luogu P2014 選課

題目大意：共有 \(n\) 門課，每門課有不同的學分。每門課沒有或有唯一一門直接的先修課程。問在修 \(m\) 門課的前提下，能夠獲得的最大學分數是多少？\(n,m\leq 300\)。

Solution：

因為每門課的先修課最多只有一門（對應着樹中每個節點至多只有 \(1\) 個父節點），所以這 \(n\) 門課程構成了森林結構（若干棵樹，因為可能有不止一門課沒有先修課）。我們可以新建一門 \(0\) 學分的課程（設這門課程編號為 \(0\)），作為“實際上沒有先修課的課程”的先修課，把包含 \(n\) 個節點的森林轉化為包含 \(n+1\) 個節點的樹，其中節點 \(0\) 為根節點。

令 \({dp}_{i,j}\) 表示在以 \(i\) 為根的子樹中選 \(j\) 門課能夠獲得的最高學分。修完 \(u\) 這門課后，對於所有的 \(v_i\in son(u)\)，我們可以在以 \(v_i\) 為根的子樹中選修若干門課（記為 \(c_i\)），在滿足 \(\sum c_i=t-1\) 的基礎上獲得盡量多的學分。

首先，顯然有 \({dp}_{u,0}=0\)。

\({dp}_{u,t}=\max\limits_{\sum\limits_{i=1}^{\left| son(u)\right|}c_i=t-1}\begin{Bmatrix}\sum\limits_{i=1}^{\left| son(u)\right|} {dp}_{v_i,c_i}\end{Bmatrix}+a_x\)

事實上，這是一個分組背包的模型。

總共有 \(\left| son(u)\right|\) 組物品，每組物品都有 \(t-1\) 個，其中第 \(i\) 組的第 \(j\) 個物品的體積為 \(j\)，價值為 \({dp}_{v_i,j}\)，背包的總容積為 \(t-1\)。我們要從每組中選出不超過 \(1\) 個物品（每個子結點 \(v\) 只能選一個狀態轉移到 \(u\)），使得物品體積不超過 \(t-1\) 的前提下（在修完 \(u\) 后，還能選修 \(t-1\) 門課），物品價值總和最大（獲得最多學分）。特別地，\(u=0\) 是一個特例，因為虛擬的根結點實際上不需要被選修，此時背包總體積應為 \(t\)。我們用分組背包進行樹形 dp 的狀態轉移。

void dfs(int x,int fa){
    f[x][0]=0;
    for(int i=hd[x];i;i=nxt[i]){    //循環子節點（物品） 
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x);
        for(int t=m;t>=0;t--)    //倒序循環當前選課總門數（當前背包體積） 
            for(int j=t;j>=0;j--)    //循環更深子樹上的選課門數（組內物品）。此處使用倒序是為了正確處理組內體積為 0 的物品
                if(t-j>=0) f[x][t]=max(f[x][t],f[x][t-j]+f[y][j]);
    }
    if(x!=0) for(int t=m;t>0;t--) f[x][t]=f[x][t-1]+a[x];    //x 不為 0 時，選修 x 本身需要占用 1 門課，獲得相應學分 
}

這類題目被稱為背包類樹形 dp，它實際上是背包與樹形 dp 的結合。除了以“節點編號”作為樹形 dp 的階段，通常我們也像線性 dp 一樣，把當前背包的體積作為第二維狀態。在狀態轉移時，我們要處理的實際上就是一個分組背包的問題。

四、換根 DP

給定一個樹形結構，需要以 每個節點為根 進行一系列統計。

考慮朴素的解法：枚舉每個節點，計算以它為根的答案。顯然復雜度不夠優秀。

我們一般通過兩次掃描來求解此類題目：

• 1. 第一次掃描時，任選一個點為根，在“有根樹”上執行一次 樹形 DP，也就是在回溯時發生的、自底向上的狀態轉移。

• 2. 第二次掃描時，從剛才選出的根出發，對整棵樹執行一次 深度優先遍歷，在每次遞歸前進行自頂向下的推導，計算出“換根”后的解。

換言之，假設當前的根是當前節點的父親，我們下一步需要將根換成當前節點。這樣就可以一直做下去。具體來說，我們需要做兩件事：

• 1. 把當前節點對父親的貢獻，從父親的 dp 值里扣除（但不能直接修改，因為父親還有別的兒子，所以最好做個備份）。

• 2. 把父親（除去當前節點的貢獻以后，剩余的部分）作為一個新的兒子，加入到當前節點的 dp 值中。這個是要直接修改的，因為要把當前節點換成根。

五、例題

1. HDU 6035 Colorful Tree

題目大意：給出一棵 \(n\) 個節點的樹，每個節點擁有一個顏色 \(c_i\)，現在定義兩點間的距離為其路徑上出現過的不同顏色數量。求兩兩點對距離之和。\(n\leq 2000\)。

Solution：

我們可以考慮每種顏色，統計經過該種顏色的路徑條數。

補集轉化，統計不經過該種顏色的路徑條數。

可以想象成是將該種顏色的點在圖中刪去，剩下的每個連通塊塊內的路徑數和就是答案。我們只要知道連通塊的大小就可以求出相應的路徑條數。

舉個栗子，如圖所示，樹上所有的粉色節點將整棵樹分為了 \(5\) 個連通塊（已在圖中用數字標出）。

考慮顏色 \(c\)，它會把樹分成很多個連通塊，每個連通塊會有 \(C_{size}^2\) 的貢獻。以 \(1\) 號點為根，與根相連的連通塊最后再特殊考慮，其他的連通塊頂端會連着一個顏色為 \(c\) 的點，在這個點處計算這個連通塊的大小，設這個點為 \(u\)。

記 \({sum}_u\) 表示 \(u\) 的子樹中到 \(u\) 的路徑上不存在其他顏色為 \(c\) 的點的個數。對 \(u\) 的每個兒子 \(v\)，我們要算出 \(v\) 的子樹中所有顏色為 \(c\) 的點的 \(sum\) 的和 \(S\)，\(sz_v-S\) 即為這個連通塊的大小。利用 \(S\) 我們也可以求出 \({sum}_u\)。

#include<bits/stdc++.h>
#define int long long
#define MEM(x,y) memset(x,y,sizeof(x))
using namespace std;
const int N=2e5+5;
int t,n,m,x,y,c[N],tot,cnt,hd[N],to[N<<1],nxt[N<<1],sz[N],k[N],sum,v,ans;
bool vis[N];
void add(int x,int y){
    to[++cnt]=y,nxt[cnt]=hd[x],hd[x]=cnt;
}
void dfs(int x,int fa){
    sz[x]=1,k[c[x]]++;    //sz[x]:以 x 為根的子樹大小 
    int p=k[c[x]];    //原來與 c[x] 有關的節點數
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x),sz[x]+=sz[y],v=sz[y]-(k[c[x]]-p);    //v:當前子樹中對應的連通量 
        sum+=v*(v-1)/2,p=(k[c[x]]+=v);    //C(v,2)=v*(v-1)/2 
    }
}
signed main(){
    while(~scanf("%lld",&n)){ 
        MEM(vis,0),MEM(hd,0),MEM(k,0),MEM(sz,0),cnt=tot=sum=ans=0;
        for(int i=1;i<=n;i++){
            scanf("%lld",&c[i]);
            if(!vis[c[i]]) tot++,vis[c[i]]=1;    //tot:顏色總數 
        }
        for(int i=1;i<n;i++){
            scanf("%lld%lld",&x,&y);
            add(x,y),add(y,x);
        }
        dfs(1,0),ans=tot*n*(n-1)/2-sum;
        for(int i=1;i<=n;i++)
            if(vis[i]) v=n-k[i],ans-=v*(v-1)/2;
        printf("Case #%lld: %lld\n",++t,ans);
    } 
    return 0;
}

2. CF1101D GCD Counting

題目大意：給出—棵 \(n\) 個節點的樹，每個節點上有點權 \(a_i\)。求最長的樹上路徑，滿足條件：路徑上經過節點（包括兩個端點）點權的 \(\gcd\) 不等於 \(1\)。\(n\leq 2\times 10^5,1\leq a_i\leq 2\times 10^5\)。

Solution：

\(\gcd=d\neq 1\)，那么肯定存在一個質數 \(p\) 滿足 \(p\mid d\)（即這條合法的鏈上的每個節點的點權都能被 \(p\) 整除）。

令 \({dp}_{i,p}\) 表示以 \(i\) 為根的子樹中能被 \(p\) 整除的最長鏈。

\(2\times 3\times 5\times 7\times 11\times 13=30030>2\times 10^5\)，所以 \(dp\) 數組的第二維開 \(6\) 就足夠了。

只需要考慮以一個點為根的子樹中，能夠整除根的點權的質因子。

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=2e5+5;
int n,a[N],x,y,cnt,hd[N],to[N<<1],nxt[N<<1],dp[N][6],ans;
vector<int>p[N]; 
void add(int x,int y){
    to[++cnt]=y,nxt[cnt]=hd[x],hd[x]=cnt;
}
void dfs(int x,int fa){
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x);
        for(int j=0;j<p[x].size();j++)    //枚舉父親的質因子
            for(int k=0;k<p[y].size();k++){    //枚舉兒子的質因子
                if(p[x][j]!=p[y][k]) continue;    //如果兩者不相等則跳過
                ans=max(ans,dp[x][j]+dp[y][k]);
                dp[x][j]=max(dp[x][j],dp[y][k]+1);     //轉移
            }
    }
}
void solve(int x,int num){    //預處理每個點的質因子 
    int cnt=0;
    for(int i=2;i<=sqrt(x);i++){
        if(x%i!=0) continue;
        p[num].push_back(i),dp[num][cnt++]=1;
        while(x%i==0) x/=i;
    } 
    if(x!=1) p[num].push_back(x),dp[num][cnt++]=1;
}
signed main(){
    scanf("%lld",&n);
    for(int i=1;i<=n;i++){ 
        scanf("%lld",&a[i]),solve(a[i],i);
        if(a[i]!=1) ans=1;
    } 
    for(int i=1;i<n;i++){
        scanf("%lld%lld",&x,&y);
        add(x,y),add(y,x);
    }
    dfs(1,0),printf("%lld\n",ans);
    return 0;
}

3. Luogu P3177「HAOI 2015」樹上染色

題目大意：有一棵點數為 \(n\) 的樹，樹邊有邊權。給你一個在 \(0 \sim n\) 之內的正整數 \(k\) ，你要在這棵樹中選擇 \(k\) 個點，將其染成黑色，並將其他的 \(n−k\) 個點染成白色。將所有點染色后，你會獲得黑點兩兩之間的距離加上白點兩兩之間的距離的和的受益。問受益最大值是多少。\(n,k\leq 2000\)。

Solution：

考慮每條邊對答案的貢獻。即，邊一側的黑點數 \(\times\) 另一側的黑點數 \(\times\) 邊權 \(+\) 一側的白點數 \(\times\) 另一側的白點數 \(\times\) 邊權。

令 \({dp}_{u,t}\) 表示以 \(u\) 為根的子樹中，有 \(t\) 個點被染成了黑色對答案貢獻的最大值。

轉化為了樹形背包問題。

枚舉更深子樹上選擇的黑點個數 \(j\)。\({dp}_{u,t}=\max({dp}_{u,t},{dp}_{u,t-j}+{dp}_{v,j}+val)\)。

邊 \((u,v)\) 對答案的貢獻 \(val\)：\(val=j\times (k-j)\times w+(sz_v-j)\times (n-k-(sz_v-j))\times w\)。

說明：\(w\) 為 \((u,v)\) 的邊權。\(k\) 為總黑點數，\(j\) 為邊一側的黑點數，那么邊另一側的黑點數就是 \(k-j\)。\(sz_v\) 表示 \(v\) 的子樹大小，那么 \(sz_v-j\) 就是邊一側的白點數。\(n-k\) 為總白點數，則另一側的白點數為 \(n-k-(sz_v-j)\)。

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=2e3+5; 
int n,k,x,y,z,cnt,hd[N],to[N<<1],nxt[N<<1],w[N<<1],sz[N],f[N][N];
void add(int x,int y,int z){
    to[++cnt]=y,nxt[cnt]=hd[x],hd[x]=cnt,w[cnt]=z;
}
void dfs(int x,int fa){
    sz[x]=1,f[x][0]=f[x][1]=0;    //不選和只選一個一定合法，故把值賦為 0  
    for(int i=hd[x];i;i=nxt[i]){
        int y=to[i];
        if(y==fa) continue;
        dfs(y,x),sz[x]+=sz[y];
        for(int t=min(k,sz[x]);t>=0;t--)    //枚舉當前黑點數 
            for(int j=0;j<=min(t,sz[y]);j++){    //枚舉更深子樹上的黑點數 
                if(f[x][t-j]==-1) continue;    //不合法則跳過 
                int val=j*(k-j)*w[i]+(sz[y]-j)*(n-k-(sz[y]-j))*w[i];    //val 
                f[x][t]=max(f[x][t],f[x][t-j]+f[y][j]+val);    //轉移 
            }
    }
}
signed main(){
    memset(f,-1,sizeof(f));
    scanf("%lld%lld",&n,&k);
    for(int i=1;i<n;i++){ 
        scanf("%lld%lld%lld",&x,&y,&z);
        add(x,y,z),add(y,x,z);
    }
    dfs(1,0),printf("%lld\n",f[1][k]);
    return 0;
} 

六、習題

• HDU6201 transaction transaction transaction
• HDU2196  Computer
• UVA10859 放置街燈 Placing Lampposts
• Luogu P4827 Crash 的文明世界（第二類斯特林數+換根 dp）