LeetCode 834 樹中距離之和

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

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

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

參考標準解析：連結寫的答案。

from typing import *


class Solution:
    def __init__(self):
        self.ans = None
        self.num_nodes = None
        self.dp = None
        self.graph = None
        self.visited = None

    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        self.ans = [0] * N
        self.num_nodes = [0] * N
        self.dp = [0] * N
        self.visited = [False] * N
        self.graph = [[] for _ in range(N)]
        for edge in edges:
            self.graph[edge[0]].append(edge[1])
            self.graph[edge[1]].append(edge[0])
        self.dfs1(0)
        self.visited = [False] * N
        self.dfs2(0)
        return self.ans

    def dfs1(self, i):
        self.visited[i] = True
        self.num_nodes[i] = 1
        for j in self.graph[i]:
            if not self.visited[j]:
                self.dfs1(j)
                self.dp[i] += (self.dp[j] + self.num_nodes[j])
                self.num_nodes[i] += self.num_nodes[j]

    def dfs2(self, i):
        self.ans[i] = self.dp[i]
        self.visited[i] = True
        for j in self.graph[i]:
            if not self.visited[j]:
                self.dp[i] -= (self.dp[j] + self.num_nodes[j])
                self.num_nodes[i] -= self.num_nodes[j]
                self.dp[j] += (self.dp[i] + self.num_nodes[i])
                self.num_nodes[j] += self.num_nodes[i]
                self.dfs2(j)
                self.num_nodes[j] -= self.num_nodes[i]
                self.dp[j] -= (self.dp[i] + self.num_nodes[i])
                self.num_nodes[i] += self.num_nodes[j]
                self.dp[i] += (self.dp[j] + self.num_nodes[j])