Leetcode834. 树中距离之和

Every day a Leetcode

题目来源：834. 树中距离之和

解法1：换根 DP

题解：【图解】一张图秒懂换根 DP！（Python/Java/C++/Go/JS）

暴力做法是，以点 i 为为树根，从 i 出发对树进行深度优先搜索，那么 i 到 j 的距离就是 j 在这棵树的深度，所有点的深度之和就是 answer[i]。

但这样做，DFS 一次的时间是 O(n)，n 个节点个点各 DFS 一次，总时间就是 O(n²)，会超时。如何优化呢？

子树大小是怎么算的？

先说二叉树，子树 x 的大小等于左子树的大小，加上右子树的大小，再加上 1（节点 x 本身），那么后序遍历这棵树，就可以算出每棵子树的大小。然后推广到一般树，子树 x 的大小，等于 x 的所有儿子的子树大小之和，再加上 1（节点 x 本身）。

问：这种算法的本质是什么？

答：以图中的这棵树为例，从「以 0 为根」换到「以 2 为根」时，原来 2 的子节点还是 2 的子节点，原来 1 的子节点还是 1 的子节点，唯一改变的是 0 和 2 的父子关系。由此可见，一对节点的距离的「变化量」应该是很小的，那么找出「变化量」的规律，就可以基于 answer[0] 算出 answer[2] 了。这种算法叫做换根 DP。

代码：

/*
 * @lc app=leetcode.cn id=834 lang=cpp
 *
 * [834] 树中距离之和
 */

// @lc code=start

// 换根 DP（树型 DP 的变种）

class Solution
{
public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges)
    {
        // g[x] 表示 x 的所有邻居
        vector<vector<int>> g(n);
        for (auto &edge : edges)
        {
            int x = edge[0], y = edge[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }
        // size[i] 为 子树 i 的大小
        vector<int> size(n, 1); // 注意这里初始化成 1 了，下面只需要累加儿子的子树大小
        vector<int> answer(n);

        function<void(int, int, int)> dfs = [&](int x, int father, int depth)
        {
            answer[0] += depth; // depth 为 0 到 x 的距离
            // 遍历 x 的邻居 y
            for (int &y : g[x])
                if (y != father) // 避免访问父节点
                {
                    dfs(y, x, depth + 1); // x 是 y 的父节点
                    size[x] += size[y];   // 累加 x 的儿子 y 的子树大小
                }
        };

        dfs(0, -1, 0); // 0 没有父节点

        function<void(int, int)> reRoot = [&](int x, int father)
        {
            // 遍历 x 的邻居 y
            for (int y : g[x])
            {
                if (y != father) // 避免访问父节点
                {
                    answer[y] = answer[x] + n - 2 * size[y];
                    reRoot(y, x); // x 是 y 的父节点
                }
            }
        };

        reRoot(0, -1); // 0 没有父节点

        return answer;
    }
};
// @lc code=end

结果：

复杂度分析：

时间复杂度：O(n)，其中 n 是树的节点个数。DFS 两次，每次 DFS 会递归访问每个节点恰好一次，所以时间复杂度为 O(n)。

空间复杂度：O(n)，其中 n 是树的节点个数。

拓展题：2858. 可以到达每一个节点的最少边反转次数

题目链接：2858. 可以到达每一个节点的最少边反转次数

代码：

/*
 * @lc app=leetcode.cn id=2858 lang=cpp
 *
 * [2858] 可以到达每一个节点的最少边反转次数
 */

// @lc code=start
class Solution
{
public:
    vector<int> minEdgeReversals(int n, vector<vector<int>> &edges)
    {
        // g[x] 表示 x 的所有邻居
        vector<vector<pair<int, int>>> g(n);
        for (auto &edge : edges)
        {
            int x = edge[0], y = edge[1];
            // x->y 视为正向，标记为 1；反向标记为 -1
            g[x].push_back(pair<int, int>(y, 1));
            g[y].push_back(pair<int, int>(x, -1));
        }
        // size[i] 为 子树 i 的大小
        // vector<int> size(n, 1); // 注意这里初始化成 1 了，下面只需要累加儿子的子树大小
        vector<int> answer(n);

        function<void(int, int)> dfs = [&](int x, int father)
        {
            // 遍历 x 的邻居 y
            for (auto &[y, dir] : g[x])
                if (y != father) // 避免访问父节点
                {
                    answer[0] += dir < 0;
                    dfs(y, x); // x 是 y 的父节点
                }
        };

        dfs(0, -1); // 0 没有父节点

        function<void(int, int)> reRoot = [&](int x, int father)
        {
            // 遍历 x 的邻居 y
            for (auto &[y, dir] : g[x])
                if (y != father) // 避免访问父节点
                {
                    answer[y] = answer[x] + dir;
                    reRoot(y, x); // x 是 y 的父节点
                }
        };

        reRoot(0, -1); // 0 没有父节点

        return answer;
    }
};
// @lc code=end

复杂度分析：

时间复杂度：O(n)，其中 n 是树的节点个数。DFS 两次，每次 DFS 会递归访问每个节点恰好一次，所以时间复杂度为 O(n)。

空间复杂度：O(n)，其中 n 是树的节点个数。