LeetCode 834. 树中距离之和

◎

本节目标

先汇总子树规模和距离，再换根计算每个节点的距离和。

题意与约束

给一棵无向树，返回每个节点到所有节点距离的总和。

第一反应与重复子问题

逐点 BFS 会重复经过边。把树先根化后，相邻根的答案只相差子树内外节点各移动一步。

状态定义与转移推导

后序求 size[node] 和根 0 的总距离。换根到子节点时：answer[child] = answer[parent] + n - 2 * size[child]。

正确性依据

换根后子树内 size[child] 个节点距离各减一，外部 n-size[child] 个节点各加一，差即为公式。

样例执行过程

官方六节点树先得到根 0 的距离和 8，沿边下传后得到 [8,12,6,10,10,10]。

代码实现

C++

C++17

#include <iostream>
#include <vector>
using namespace std;

class Solution {
public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        vector<vector<int>> graph(n);
        for (const auto& edge : edges) {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }
        vector<int> parent(n, -1);
        vector<int> order{0};
        vector<int> size(n, 1);
        vector<int> answer(n);
        for (int index = 0; index < static_cast<int>(order.size()); index++) {
            int node = order[index];
            for (int nxt : graph[node]) {
                if (nxt != parent[node]) {
                    parent[nxt] = node;
                    order.push_back(nxt);
                }
            }
        }
        for (int index = n - 1; index > 0; index--) {
            int node = order[index];
            size[parent[node]] += size[node];
            answer[0] += size[node];
        }
        for (int node : order) {
            for (int nxt : graph[node]) {
                if (parent[nxt] == node) {
                    answer[nxt] = answer[node] + n - 2 * size[nxt];
                }
            }
        }
        return answer;
    }
};

Python

Python 3

class Solution:
    def sumOfDistancesInTree(self, n, edges):
        graph = [[] for _ in range(n)]
        for left, right in edges:
            graph[left].append(right)
            graph[right].append(left)
        parent, order = [-1] * n, [0]
        for node in order:
            for nxt in graph[node]:
                if nxt != parent[node]:
                    parent[nxt] = node
                    order.append(nxt)
        size, answer = [1] * n, [0] * n
        for node in reversed(order[1:]):
            size[parent[node]] += size[node]
            answer[0] += size[node]
        for node in order:
            for nxt in graph[node]:
                if parent[nxt] == node:
                    answer[nxt] = answer[node] + n - 2 * size[nxt]
        return answer

复杂度分析

两遍均线性，时间 O(n)，邻接表、父序和数组空间 O(n)。

边界与易错点

单节点答案为 [0]；Python 使用父序正向建树、逆序聚合，避免递归深度风险。

模式迁移

“所有节点作为根”的树题优先寻找从父答案到子答案的增量公式。回到树形动态规划。