今日份题目:
给定一个无向、连通的树。树中有 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;
}
};
提交结果