給定一個無向、連通的樹。樹中有 N 個標記為 0...N-1 的節點以及 N-1?條邊?。

第 i 條邊連接節點?edges[i][0] 和 edges[i][1]?。

返回一個表示節點 i 與其他所有節點距離之和的列表 ans。

示例 1:

輸入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]] 輸出: [8,12,6,10,10,10] 解釋: 如下為給定的樹的示意圖：0/ \ 1 2/|\3 4 5我們可以計算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此類推。

說明:?1 <= N <= 10000

樹形動態規劃

首先我們來考慮一個節點的情況，即每次題目指定一棵樹，以 $\text{root}$ 為根，詢問節點 $\text{root}$ 與其他所有節點的距離之和。

很容易想到一個樹形動態規劃：定義 $\text{dp}[u]$ 表示以 $u$ 為根的子樹，它的所有子節點到它的距離之和，同時定義 $\text{sz}[u]$ 表示以 $u$ 為根的子樹的節點數量，不難得出如下的轉移方程：

dp[u]=∑v∈son[u]dp[v]+sz[v]\textit{dp}[u]=\sum_{v\in \textit{son}[u]}\textit{dp}[v] + \textit{sz}[v] dp[u]=v∈son[u]∑?dp[v]+sz[v]

其中 $\text{son}[u]$ 表示 $u$ 的所有后代節點集合。轉移方程表示的含義就是考慮每個后代節點 $v$，已知 $v$ 的所有子節點到它的距離之和為 $\text{dp}[v]$，那么這些節點到 $u$ 的距離之和還要考慮 $u\to v$ 這條邊的貢獻。考慮這條邊長度為 $1$，一共有 $sz[v]$ 個節點到節點 $u$ 的距離會包含這條邊，因此貢獻即為 $1\times \text{sz}[v]=\text{sz}[v]$。我們遍歷整棵樹，從葉子節點開始自底向上遞推到根節點 $\text{root}$ 即能得出最后的答案為 $\text{dp}[\text{root}]$。

回到本題中，題目要求的其實是上題的擴展，即要求我們求出每個節點為根節點的時候，它與其他所有節點的距離之和。暴力的角度我們可以考慮對每個節點都做一次如上的樹形動態規劃，這樣時間復雜度即為 O(N2)O(N^2)O(N2)，那么有沒有更優雅的方法呢？

經過一次樹形動態規劃后其實我們獲得了在 $u$ 為根的樹中，每個節點為根的子樹的答案 $\text{dp}$，我們可以利用這些已有信息來優化時間復雜度。

假設 $u$ 的某個后代節點為 $v$，如果要算 $v$ 的答案，本來我們要以 $v$ 為根來進行一次樹形動態規劃。但是利用已有的信息，我們可以考慮樹的形態做一次改變，讓 $v$ 換到根的位置，$u$ 變為其孩子節點，同時維護原有的 $\text{dp}$ 信息。在這一次的轉變中，我們觀察到除了 $u$ 和 $v$ 的 $\text{dp}$ 值，其他節點的 $\text{dp}$ 值都不會改變，因此只要更新 $\text{dp}[u]$ 和 $\text{dp}[v]$ 的值即可。

那么我們來看 $v$ 換到根的位置的時候怎么利用已有信息求出 $\text{dp}[u]$ 和 $\text{dp}[v]$ 的值。重新回顧第一次樹形動態規劃的轉移方程，我們可以知道當 $u$ 變為 $v$ 的孩子的時候 $v$ 不在 $u$ 的后代集合 $\text{son}[u]$ 中了，因此此時 $\text{dp}[u]$ 需要減去 $v$ 的貢獻，即

$$\text{dp}[u]=\text{dp}[u]?(\text{dp}[v]+\text{sz}[v])$$

同時 $\text{sz}[u]$ 也要相應減去 $\text{sz}[v]$。

而 $v$ 的后代節點集合中多出了 $u$，因此 $\text{dp}[v]$ 的值要由 $u$ 更新上來，即

$$\text{dp}[v]=\text{dp}[v]+(\text{dp}[u]+\text{sz}[u])$$

同時 $\text{sz}[v]$ 也要相應加上 $\text{sz}[u]$。

至此我們完成了一次「換根」操作，在 $O(1)$ 的時間內維護了 $\text{dp}$ 的信息，且此時的樹結構以 $v$ 為根。那么接下來我們不斷地進行換根的操作，即能在 $O(N)$ 的時間內求出每個節點為根的答案，實現了時間復雜度的優化。

Code

def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:def dfs1(u, f):sz[u] = 1dp[u] = 0for v in graph[u]:if v == f:continuedfs1(v, u)dp[u] += (dp[v] + sz[v])sz[u] += sz[v]def dfs2(u, f):ans[u] = dp[u]for v in graph[u]:if v == f:continuepu = dp[u]pv = dp[v]su = sz[u]sv = sz[v]dp[u] -= (dp[v] + sz[v])sz[u] -= sz[v]dp[v] += (dp[u] + sz[u])sz[v] += sz[u]dfs2(v, u)dp[u] = pudp[v] = pvsz[u] = susz[v] = svsz = [0 for _ in range(N)]dp = [0 for _ in range(N)]ans = [0 for _ in range(N)]graph = [[] for _ in range(N)]for edge in edges:u = edge[0]v = edge[1]graph[u].append(v)graph[v].append(u)dfs1(0, -1)dfs2(0, -1)return ans

總結

以上是生活随笔為你收集整理的834. Sum of Distances in Tree的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。