如何优化C++求解树中节点距离和的代码以解决TLE超时问题

你当前代码超时的核心原因有两点：

• 时间复杂度为O(n²)：针对每个节点单独跑一次BFS遍历全树，当n规模达到1e4及以上时完全无法通过时间限制
• 额外性能损耗：使用map<pair<int,int>>存储所有两两节点距离，插入和查询都带有*O(logn)*的额外开销，完全没有必要——就算保留BFS思路，也可以在遍历过程中直接把距离累加到结果数组，不需要额外存储所有配对距离。

即便优化掉map的问题，O(n²)的复杂度还是无法满足大数据量的测试用例，最优的解法是使用换根树形DP，时间复杂度可以降到O(n)。

只需要两次DFS遍历全树即可得到所有节点的距离和，核心思路如下：

1. 第一次后序遍历（默认以0为根节点）：
   • 计算cnt[u]：以u为根的子树包含的总节点数
   • 计算res[u]：以u为根的子树中，所有节点到u的距离之和
2. 第二次前序遍历（换根推导）：
   当根节点从父节点u切换到它的子节点v时，所有节点的距离和可以通过父节点的结果直接推导：
   • v的子树内共有cnt[v]个节点，到v的距离比到u的距离少1，总贡献减少cnt[v]
   • 剩余n - cnt[v]个节点不在v的子树内，到v的距离比到u的距离多1，总贡献增加n - cnt[v]
     因此推导公式为：res[v] = res[u] - cnt[v] + (n - cnt[v])

class Solution {
public:
    vector<vector<int>> g;
    vector<int> res;
    vector<int> cnt;
    int n;

    void dfs1(int u, int parent) {
        cnt[u] = 1;
        res[u] = 0;
        for (int v : g[u]) {
            if (v == parent) continue;
            dfs1(v, u);
            cnt[u] += cnt[v];
            res[u] += res[v] + cnt[v];
        }
    }

    void dfs2(int u, int parent) {
        for (int v : g[u]) {
            if (v == parent) continue;
            res[v] = res[u] - cnt[v] + n - cnt[v];
            dfs2(v, u);
        }
    }

    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        this->n = n;
        g.resize(n);
        res.resize(n);
        cnt.resize(n);
        for (auto& e : edges) {
            g[e[0]].push_back(e[1]);
            g[e[1]].push_back(e[0]);
        }
        dfs1(0, -1);
        dfs2(0, -1);
        return res;
    }
};

针对你给出的输入n=6，边为[[0,1],[0,2],[2,3],[2,4],[2,5]]：

1. 第一次dfs1计算得到res[0]=8，和示例给出的结果一致
2. 第二次换根计算得到res[2] = 8 -4 + 2 =6、res[1]=8-1+5=12，完全匹配示例输出的[8,12,6,10,10,10]

内容的提问来源于stack exchange，提问作者Rohit gupta