每天一道leetcode：834.树中距离之和

最新推荐文章于 2024-08-10 09:23:37 发布

杜若南星

最新推荐文章于 2024-08-10 09:23:37 发布

阅读量192

点赞数 1

文章标签： leetcode 算法职场和发展数据结构深度优先

本文链接：https://blog.csdn.net/qq_63349644/article/details/131747641

版权

该文章描述了一个关于无向、连通树的算法问题，其中目标是计算每个节点到所有其他节点的距离之和。解决方案使用了动态规划和深度优先遍历，通过两次DFS遍历更新节点的距离和子树大小，最终得到每个节点的距离总和。

摘要由CSDN通过智能技术生成

今日份题目：

给定一个无向、连通的树。树中有 n 个标记为 0...n-1 的节点以及 n-1 条边。

给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。

示例1

输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 树如图所示。
我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

示例2

输入: n = 1, edges = []
输出: [0]

示例3

输入: n = 2, edges = [[1,0]]
输出: [1,1]

提示

1 <= n <= 3 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
给定的输入保证为有效的树

题目思路

动态规划，定义 dp[u] 表示以 u 为根的子树，它的所有子节点到它的距离之和，同时定义 sz[u] 表示以 u 为根的子树的节点数量，son[u] 表示 u 的所有后代节点集合

状态转移方程：dp[u]=v∈son[u]∑dp[v]+sz[v]

代码

这道题目的难度是困难，确实是不是很简单，我就学习了一下官方答案，添加注释如下：

class Solution 
{
public:
    vector<int> ans, sz, dp;//结果数组、子树个数数组、所有子树距离和数组
    vector<vector<int>> graph;//图

    void dfs(int u, int f) 
    {
        sz[u] = 1;//初始化，开始时只有自己，为1
        dp[u] = 0;//自己到自己的距离为0
        for (auto& v: graph[u]) //遍历图中与u相连的点
        {
            if (v == f) //无效数据
            {
                continue;
            }
            dfs(v, u);//递归遍历下一个子节点
            dp[u] += dp[v] + sz[v];//根据状态转移方程更新
            sz[u] += sz[v];//v的子树个数并入u的子树个数中
        }
    }

    void dfs2(int u, int f) 
    {
        ans[u] = dp[u];//记录节点u的答案
        for (auto& v: graph[u]) //遍历图中与u相连的点
        {
            if (v == f) //无效数据
            {
                continue;
            }
            //转换为从当前节点到其他节点
            int pu = dp[u], pv = dp[v];
            int su = sz[u], sv = sz[v];
			//把树翻过来，出发的节点减去，加到到达的节点
            dp[u] -= dp[v] + sz[v];
            sz[u] -= sz[v];
            dp[v] += dp[u] + sz[u];
            sz[v] += sz[u];

            dfs2(v, u);//反向深度优先遍历

            dp[u] = pu, dp[v] = pv;//恢复初始信息
            sz[u] = su, sz[v] = sv;
        }
    }

    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) 
    {
        ans.resize(n, 0);//设置vector数组长度为n，初始化为0
        sz.resize(n, 0);
        dp.resize(n, 0);
        graph.resize(n, {});
        for (auto& edge: edges) //存入图信息
        {
            int u = edge[0], v = edge[1];
            graph[u].emplace_back(v);
            graph[v].emplace_back(u);
        }
        dfs(0, -1);
        dfs2(0, -1);
        return ans;
    }
};