---
title: "LeetCode 834 (換根DP)"
tags: 題解, 樹, 動態規劃
---

https://leetcode.com/problems/sum-of-distances-in-tree/

#### 題意

給一棵無根樹，求出以每個節點作為根時，所有點與根的距離之和。

#### 思路

換根DP，先 DFS 求出以 $u$ 為根時，子樹的節點數以及節點到 $u$ 的距離之和。

第二次 DFS 考慮 $u$ 為根而 $v$ 為其子節點時，已知 $dist[u]$ 為所有點到 $u$ 的距離之和，而 $dist[v]$ 為以 $v$ 為根的子樹中所有點到 $v$ 的距離之和。也就是說，對於 $v$ 而言，還差以 $u$ 為根的子樹的貢獻，將 $dist[u]$ 排除 $v$ 子樹的貢獻後，加上 $u$ 子樹的節點數，即為從 $u$ 轉移至 $v$ 的貢獻。

#### Code

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        def dfs(node, prev):
            # get sum of dist and num of nodes from
            # subtree rooted at `node`
            for neigh in adj[node]:
                if neigh == prev:
                    continue
                dfs(neigh, node)
                dists[node] += dists[neigh] + nodes[neigh]
                nodes[node] += nodes[neigh]

        def dfs2(node, prev):
            # rerooting dp
            for neigh in adj[node]:
                if neigh == prev:
                    continue
                dists[neigh] += (dists[node] - (dists[neigh] + nodes[neigh])) + (n - nodes[neigh])
                dfs2(neigh, node)

        adj = defaultdict(list)
        for (a, b) in edges:
            adj[a].append(b)
            adj[b].append(a)

        dists = [0 for _ in range(n)]
        nodes = [1 for _ in range(n)]
        dfs(0, -1)
        dfs2(0, -1)
        return dists