C++ 中树的距离之和
假设我们有一个无向连通树,其中存在 N 个节点。这些节点标记为 0...N-1,并且给出 N-1 条边。第 i 条边连接节点 edges[i][0] 和 edges[i][1]。我们必须找到一个列表,其中 ans[i] 是节点 i 与所有其他节点之间距离的总和。
因此,如果输入类似于 N = 6 且 edges = [(0,1),(0,2),(2,3),(2,4),(2,5)],则输出将为 [8,12,6,10,10,10]
为了解决这个问题,我们将遵循以下步骤:
定义一个函数 dfs1(),它将接收节点、父节点作为参数。
对于初始化 i := 0,当 i < graph[node] 的大小,更新(i 增加 1),执行:
child := graph[node, i]
如果 child 不等于 parent,则:
dfs1(child, node)
cnt[node] := cnt[node] + cnt[child]
ans[node] := ans[node] + cnt[child] + ans[child]
定义一个函数 dfs2(),它将接收节点、父节点作为参数。
对于初始化 i := 0,当 i < graph[node] 的大小,更新(i 增加 1),执行:
child := graph[node, i]
如果 child 不等于 parent,则:
ans[child] := ans[node] - cnt[child] + N - cnt[child]
dfs2(child, node)
定义一个数组 ans
定义一个数组 cnt
定义一个具有 10005 行的数组 graph
从主方法中,执行以下操作:
N of this := N
ans := 定义一个大小为 N 的数组
cnt := 定义一个大小为 N 的数组,并用 1 填充它
n := edges 的大小
对于初始化 i := 0,当 i < n,更新(i 增加 1),执行:
u := edges[i, 0]
v := edges[i, 1]
在 graph[u] 的末尾插入 v
在 graph[v] 的末尾插入 u
dfs1(0, -1)
dfs2(0, -1)
返回 ans
让我们看看以下实现以获得更好的理解:
示例
#include <bits/stdc++.h> using namespace std; void print_vector(vector<auto> v){ cout << "["; for(int i = 0; i<v.size(); i++){ cout << v[i] << ", "; } cout << "]"<<endl; } class Solution { public: void dfs1(int node, int parent) { for (int i = 0; i < graph[node].size(); i++) { int child = graph[node][i]; if (child != parent) { dfs1(child, node); cnt[node] += cnt[child]; ans[node] += cnt[child] + ans[child]; } } } void dfs2(int node, int parent) { for (int i = 0; i < graph[node].size(); i++) { int child = graph[node][i]; if (child != parent) { ans[child] = ans[node] - cnt[child] + N - cnt[child]; dfs2(child, node); } } } vector<int> ans; vector<int> cnt; vector<int> graph[10005]; int N; vector<int> sumOfDistancesInTree(int N, vector<vector<int> >& edges) { this->N = N; ans = vector<int>(N); cnt = vector<int>(N, 1); int n = edges.size(); for (int i = 0; i < n; i++) { int u = edges[i][0]; int v = edges[i][1]; graph[u].push_back(v); graph[v].push_back(u); } dfs1(0, -1); dfs2(0, -1); return ans; } }; main(){ Solution ob; vector<vector<int>> v = {{0,1},{0,2},{2,3},{2,4},{2,5}}; print_vector(ob.sumOfDistancesInTree(6, v)); }
输入
{{0,1},{0,2},{2,3},{2,4},{2,5}}
输出
[8, 12, 6, 10, 10, 10, ]