動態規劃入門——動態規劃與數據結構的結合，在樹上做DP

本文由TechFlow原創，本博文僅作為知識點學習，不會用於任何商業用途！

今天我們來看一個有趣的問題，通過這個有趣的問題，我們來了解一下在樹形結構當中做動態規劃的方法。

這個問題題意很簡單，給定一棵樹，並不一定是二叉樹，樹上的樹枝帶有權重，可以看成是長度。要求樹上最長的鏈路的長度是多少？

比如我們隨手畫一棵樹，可能丑了點，勿怪：

如果讓我們用肉眼來看，稍微嘗試一下就能找到答案，最長的路徑應該是下圖當中紅色的這條：

但是如果讓我們用演算法來算，應該怎麼辦呢？

這道題其實有一個非常巧妙的辦法，我們先不講，先來看看動態規劃怎麼解決這個問題。

樹形DP

動態規劃並不只是可以在數組當中運行，實際上只要滿足動態規劃的狀態轉移的條件和無後效性就可以使用動態規劃，無論在什麼數據結構當中。樹上也是一樣的，明白了這點之後，就只剩下了兩個問題，第一個是狀態是什麼，第二個問題是狀態之間怎麼轉移？

在之前的背包問題當中，狀態就是背包當前用的體積，轉移呢就是我們新拿一個物品的決策。但是這一次我們要在樹上進行動態規劃，相對來說狀態和對應的轉移會隱蔽一些。沒有關係，我會從頭開始整理思路，一點一點將推導和思考的過程講解清楚。

首先，我們都知道，狀態之間轉移其實本質上是一個由局部計算整體的過程。我們通過相對容易的子狀態進行轉移，得到整體的結果。這個是動態規劃的精髓，某種程度上來說它和分治法也比較接近，都存在大問題和小問題之間邏輯上的關係。所以當我們面臨一個大問題一籌莫展的時候，可以借鑒一下分治法，思考一下從小問題入手。

所以，我們從小到大，由微觀到宏觀，來看看最簡單的情況：

這種情況很明顯，鏈路只有一條，所以長度自然是5 + 6 = 11，這顯然也是最長的長度。這種情況都沒有問題，下面我們來把情況稍微再變得複雜一些，我們在樹上多加入一層：

這張圖稍微複雜了一些，但是路徑也不難找到，應該是E-B-F-H。路徑的總長度為12：

但是如果我們變更一下路徑長度呢，比如我們把FG和FH的路徑加長，會得到什麼結果呢？

顯然這種情況下答案就變了，FGH是最長的。

舉這個例子只為了說明一個很簡單的問題，即對於一棵樹而言它上面的最長路徑並不一定經過根節點。比如剛才的例子當中，如果路徑必須要經過B的話，最長只能構造出4+2+16=22的長度，但是如果可以不用經過B的話，可以得到最長的長度是31。

得出這個結論看似好像沒有用，但其實對於我們理清思路很有幫助。既然我們不能保證最長路徑一定會經過樹根，所以我們就不能直接轉移答案。那我們應該怎麼辦呢？

回答這個問題光想是不夠的，依然需要我們來觀察問題和深入思考。

轉移過程

我們再觀察一下下面這兩張圖：

有沒有發現什麼規律？

由於我們的數據結構就是樹形的，所以這個最長路徑不管它連通的哪兩個節點，一定可以保證，它會經過某一棵子樹的根節點。不要小看這個不起眼的結論，實際上它非常重要。有了這個結論之後，我們將整條路徑在根節點處切開。

切開之後我們得到了兩條通往葉子節點的鏈路，問題來了，根節點通往葉子節點的鏈路有很多條，為什麼是這兩條呢？

很簡單，因為這兩條鏈路最長。所以這樣加起來之後就可以保證得到的鏈路最長。這兩條鏈路都是從葉子節點通往A的，所以我們得到的最長鏈路就是以A為根節點的子樹的最長路徑。

我們前面的分析說了，最長路徑是不能轉移的，但是到葉子的最長距離是可以轉移的。我們舉個例子：

F到葉子的最長距離顯然就是5和6中較大的那個，B稍微複雜一些，D和E都是葉子節點，這個容易理解。它還有一個子節點F，對於F來說它並不是葉子節點，但是我們前面算到了F到葉子節點的最長距離是6，所以B通過F到葉子節點的最長距離就是2 + 6 = 8。這樣我們就得到了狀態轉移方程，不過我們轉移的不是要求的答案而是從當前節點到葉子節點的最長距離和次長距離。

因為只有最長距離是不夠的，因為我們要將根節點的最長距離加上次長距離得到經過根節點的最長路徑，由於我們之前說過，所有的路徑必然經過某棵子樹的根節點。這個想明白了是廢話，但是這個條件的確很重要。既然所有的鏈路都至少經過某一個子樹的根節點，那麼我們算出所有子樹經過根節點的最長路徑，其中最長的那個不就是答案么？

下面我們演示一下這個過程：

上圖當中用粉色筆標出的就是轉移的過程，對於葉子節點來說最長距離和次長距離都是0，主要的轉移過程發生在中間節點上。

轉移的過程也很容易想通，對於中間節點i，我們遍歷它所有的子節點j，然後維護最大值和次大值，我們寫下狀態轉移方程：

狀態轉移想明白了，剩下的就是編碼的問題了。可能在樹上尤其是遞歸的時候做狀態轉移有些違反我們的直覺，但實際上並不難，我們寫出程式碼來看下，我們首先來看建樹的這個部分。為了簡化操作，我們可以把樹上所有的節點序號看成是int，對於每一個節點，都會有一個數組存儲所有與這個節點連接的邊，包括父親節點。

由於我們只關注樹上的鏈路的長度，並不關心樹的結構，樹建好了之後，不管以哪一個點為整體的樹根結果都是一樣的。所以我們隨便找一個節點作為整棵樹的根節點進行遞歸即可。強調一下，這個是一個很重要的性質，因為本質上來說，樹是一個無向無環全連通圖。所以不管以哪個節點為根節點都可以連通整棵子樹。

我們創建一個類來存儲節點的資訊，包括id和兩個最長以及次長的長度。我們來看下程式碼，應該比你們想的要簡單得多。

class Node(object):
    def __init__(self, id):
        self.id = id
        # 以當前節點為根節點的子樹到葉子節點的最長鏈路
        self.max1 = 0
        # 到葉子節點的次長鏈路
        self.max2 = 0
        # 與當前節點相連的邊
        self.edges = []

    # 添加新邊
    def add_edge(self, v, l):
        self.edges.append((v, l))


# 創建數組，存儲所有的節點
nodes = [Node(id) for id in range(12)]

edges = [(0, 1, 3), (0, 2, 1), (1, 3, 1), (1, 4, 4), (1, 5, 2), (5, 6, 5), (5, 7, 6), (2, 8, 7), (7, 9, 2), (7, 10, 8)]

# 創建邊
for edge in edges:
    u, v, l = edge
    nodes[u].add_edge(v, l)
    nodes[v].add_edge(u, l)

由於我們只是為了傳達思路，所以省去了許多面向對象的程式碼，但是對於我們理解題目思路來說應該是夠了。

下面，我們來看樹上做動態規劃的程式碼：

def dfs(u, f, ans):
    nodeu = nodes[u]
    # 遍歷節點u所有的邊
    for edge in nodes[u].edges:
        v, l = edge
        # 注意，這其中包括了父節點的邊
        # 所以我們要判斷v是不是父節點的id
        if v == f:
            continue
        # 遞歸，更新答案
        ans = max(ans, dfs(v, u, ans))
        nodev = nodes[v]
        # 轉移最大值和次大值
        if nodev.max1 + l > nodeu.max1:
            nodeu.max1 = nodev.max1 + l
        elif nodev.max1 + l > nodeu.max2:
            nodeu.max2 = nodev.max1 + l
    # 返回當前最優解
    return max(ans, nodeu.max1 + nodeu.max2)

看起來很複雜的樹形DP，其實程式碼也就只有十來行，是不是簡單得有些出人意料呢？

但是還是老生常談的話題，這十幾行程式碼看起來簡單，但是其中的細節還是有一些的，尤其是涉及到了遞歸操作。對於遞歸不是特別熟悉的同學可能會有些吃力，建議可以根據之前的圖手動在紙上驗算一下，相信會有更深刻的認識。

另一種做法

文章還沒完，我們還有一個小彩蛋。其實這道題還有另外一種做法，這種做法非常機智，也一樣介紹給大家。

之前我們說了，由於樹記錄的是節點的連通狀態，所以不管以哪個節點為根節點，都不會影響整棵樹當中路徑的長度以及結構。既然如此，如果我們富有想像力的話，我們把一棵樹壓扁，是不是可以看成是一串連在一起的繩子或者木棍？

我們來看下圖：

我們把C點向B點靠近，並不會影響樹的結構，畢竟這是一個抽象出來的架構，我們並不關注樹上樹枝之間的夾角。我們可以想像成我們拎起了A點，其他的幾點由於重力的作用下垂，最後就會被拉成一條直線。

比如上圖當中，我們拎起了A點，BCD都垂下。這個時候位於最下方的點是D點。那麼我們再拎起D點，最下方的點就成了C點，那麼DC之間的距離就是樹上的最長鏈路：

我們把整個過程梳理一下，首先我們隨便選了一個點作為樹根，然後找出了距離它最遠的點。第二次，我們選擇這個最遠的點作為樹根，再次找到最遠的點。這兩個最遠點之間的距離就是答案。

這種做法非常直觀，但是我也想不到可以嚴謹證明的方法，有思路的小夥伴可以在後台給我留言。如果有些想不通的小夥伴可以自己試著用幾根繩子連在一起，然後拎起來做個實驗。看看這樣拎兩次得到的兩個點，是不是樹上距離最遠的兩個點。

最後，我們來看下程式碼：

def dfs(u, f, dis, max_dis, nd):
    nodeu = nodes[u]
    for edge in nodes[u].edges:
        v, l = edge
        if v == f:
            continue
        nodev = nodes[v]
        # 更新最大距離，以及最大距離的點
        if dis + l > max_dis:
            max_dis, nd = dis+l, nodev
        # 遞歸
        _max, _nd = dfs(v, u, dis+l, max_dis, nd)
        # 如果遞歸得到的距離更大，則更新
        if _max > max_dis:
            max_dis, nd = _max, _nd
    # 返回
    return max_dis, nd

# 第一次遞歸，獲取距離最大的節點
_, nd = dfs(0, -1, 0, 0, None)
# 第二次遞歸，獲取最大距離
dis, _ = dfs(nd.id, -1, 0, 0, None)
print(dis)

到這裡，這道有趣的題目就算是講解完了，不知道文中的兩種做法大家都學會了嗎？第一次看可能會覺得有些蒙，問題很多這是正常的，但核心的原理並不難，畫出圖來好好演算一下，一定可以得到正確的結果。

Tags: 演算法基礎：動態規劃演算法基礎：樹形DP