834. Sum of Distances in Tree

Reference to the solution
https://leetcode.com/problems/sum-of-distances-in-tree/solution/

二刷
96.91 %
class Solution {
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        List<List<Integer>> graph = new ArrayList<>();
        for (int i = 0; i < N; i++) {
            graph.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            graph.get(edge[0]).add(edge[1]);
            graph.get(edge[1]).add(edge[0]);
        }
        int[] subCount = new int[N];
        int sum = dfs(0, graph, -1, subCount);
        int[] result = new int[N];
        result[0] = sum;
        dfs2(0, graph, N, -1, subCount, result);
        return result;
    }
   
    private void dfs2(int node, List<List<Integer>> graph, int N, int parent, int[] subCount, int[] result) {
        List<Integer> neighbors = graph.get(node);
        if (node != 0) {
            result[node] = result[parent] - subCount[node] + (N - subCount[node]);
        }
        for (int neighbor : neighbors) {
            if (neighbor == parent) {
                continue;
            }
            dfs2(neighbor, graph, N, node, subCount, result);
        }
    }
    // 1-2-3-0
   
    private int dfs(int node, List<List<Integer>> graph, int parent, int[] subCount) {
        List<Integer> neighbors = graph.get(node);
        int sum = 0;
        subCount[node] = 1;
        for (int neighbor : neighbors) {
            if (neighbor == parent) {
                continue;
            }
            sum += dfs(neighbor, graph, node, subCount) + subCount[neighbor];
            subCount[node] += subCount[neighbor];
        }
        return sum;
    }
}
一刷
Time O(n) Space O(n)

 97.62 %
class Solution {
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
List<List<Integer>> graph = new ArrayList<>();
for (int i = 0; i < N; i++) {
graph.add(new ArrayList<>());
}
for (int[] edge : edges) {
graph.get(edge[0]).add(edge[1]);
graph.get(edge[1]).add(edge[0]);
}
int[] subtreeSum = new int[N];
int[] subtreeCount = new int[N];
dfs(graph, subtreeSum, subtreeCount, 0, -1);
update(graph, subtreeSum, subtreeCount, 0, -1);
return subtreeSum;
}

private void dfs(List<List<Integer>> graph, int[] subtreeSum, int[] subtreeCount, int curr, int parent) {
List<Integer> neighbors = graph.get(curr);
subtreeCount[curr] = 1;
for (int neighbor : neighbors) {
if (neighbor != parent) {
dfs(graph, subtreeSum, subtreeCount, neighbor, curr);
subtreeSum[curr] += subtreeSum[neighbor] + subtreeCount[neighbor];
subtreeCount[curr] += subtreeCount[neighbor];
}
}
}

private void update(List<List<Integer>> graph, int[] subtreeSum, int[] subtreeCount, int curr, int parent) {
if (curr != 0) {
subtreeSum[curr] = subtreeSum[parent] - subtreeCount[curr] + graph.size() - subtreeCount[curr];
}
List<Integer> neighbors = graph.get(curr);
for (int neighbor : neighbors) {
if (neighbor != parent) {
update(graph, subtreeSum, subtreeCount, neighbor, curr);
}
}
}
}