diff options
Diffstat (limited to 'other/sum_dist_tree.cpp')
-rw-r--r-- | other/sum_dist_tree.cpp | 45 |
1 files changed, 45 insertions, 0 deletions
diff --git a/other/sum_dist_tree.cpp b/other/sum_dist_tree.cpp new file mode 100644 index 0000000..10fc449 --- /dev/null +++ b/other/sum_dist_tree.cpp @@ -0,0 +1,45 @@ +#include "leetcode.h" + +class Solution { +public: + vector<int> sumOfDistancesInTree(int n, vvd & edges) { + tree.resize(n); + ans.assign(n, 0); + count.assign(n, 1); + for (auto e : edges) { + tree[e[0]].insert(e[1]); + tree[e[1]].insert(e[0]); + } + dfs(0, -1); + dfs2(0, -1); + return ans; + } + +private: + vector<unordered_set<int>> tree; + vector<int> ans, count; + void dfs(int root, int pre) { + for (auto i : tree[root]) { + if (i == pre) + continue; + dfs(i, root); + count[root] += count[i]; + ans[root] += ans[i] + count[i]; + } + } + void dfs2(int root, int pre) { + for (auto i : tree[root]) { + if (i == pre) + continue; + ans[i] = ans[root] - count[i] + count.size() - count[i]; + dfs2(i, root); + } + } +}; + +int main() { + Solution obj; + vvd edges = {{0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5}}; + for (auto i : obj.sumOfDistancesInTree(6, edges)) + cout << i << ' '; // expect: 8 12 6 10 10 10 +} |