題目

https://leetcode.com/problems/sum-of-distances-in-tree/

題解

一般的算法題，指令條數為 10^8 以內是可以通過的。也就是說，如果 arr.length=10^6 的話，你的算法不能超過 O(n^2)。

根據本題數據范圍限制，1 <= n <= 3 * 10^4，所以常規的 Dijskra 和 Floid 在復雜度上明顯會超時。沒繼續在暴力法上耗時間，直奔答案了。

Should I quit ?
No, I will ignore this 😃
But I will be back soon to solve this question.
Congratulations to those who are able to solve this question.

Approach #1: Subtree Sum and Count [Accepted]

Intuition

Let ans be the returned answer, so that in particular ans[x] be the answer for node x.

Naively, finding each ans[x] would take O(N) time (where N is the number of nodes in the graph), which is too slow. This is the motivation to find out how ans[x] and ans[y] are related, so that we cut down on repeated work.

Let’s investigate the answers of neighboring nodes x and y. In particular, say xy is an edge of the graph, that if cut would form two trees X (containing x) and Y (containing y).

Then, as illustrated in the diagram, the answer for x in the entire tree, is the answer of x on X "x@X", plus the answer of y on Y "y@Y", plus the number of nodes in Y "#(Y)". The last part "#(Y)" is specifically because for any node z in Y, dist(x, z) = dist(y, z) + 1.

By similar reasoning, the answer for y in the entire tree is ans[y] = x@X + y@Y + #(X). Hence, for neighboring nodes x and y, ans[x] - ans[y] = #(Y) - #(X).

Algorithm

Root the tree. For each node, consider the subtree Snode of that node plus all descendants. Let count[node] be the number of nodes in Snode, and stsum[node] (“subtree sum”) be the sum of the distances from node to the nodes in Snode.

We can calculate count and stsum using a post-order traversal, where on exiting some node, the count and stsum of all descendants of this node is correct, and we now calculate count[node] += count[child] and stsum[node] += stsum[child] + count[child].

This will give us the right answer for the root: ans[root] = stsum[root].

Now, to use the insight explained previously: if we have a node parent and it’s child child, then these are neighboring nodes, and so ans[child] = ans[parent] - count[child] + (N - count[child]). This is because there are count[child] nodes that are 1 easier to get to from child than parent, and N-count[child] nodes that are 1 harder to get to from child than parent.

Using a second, pre-order traversal, we can update our answer in linear time for all of our nodes.

class Solution {int N;int[] count;int[] ans;Set<Integer>[] graph;public int[] sumOfDistancesInTree(int n, int[][] edges) {graph = new HashSet[n];for (int i = 0; i < n; i++) {graph[i] = new HashSet<>();}N = n;count = new int[n];ans = new int[n];Arrays.fill(count, 1);for (int[] pair : edges) {graph[pair[0]].add(pair[1]);graph[pair[1]].add(pair[0]);}dfs1(0, -1);dfs2(0, -1);return ans;}public void dfs1(int node, int parent) {for (int child : graph[node]) {if (child != parent) {dfs1(child, node);count[node] += count[child];ans[node] += ans[child] + count[child];}}}public void dfs2(int node, int parent) {for (int child : graph[node]) {if (child != parent) {ans[child] = ans[node] + (N - count[child]) - count[child];dfs2(child, node);}}} }

總結

以上是生活随笔為你收集整理的leetcode 834. Sum of Distances in Tree | 834. 树中距离之和（树形DP）的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。