[LeetCode] 834. Sum of Distances in Tree

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

leetcode - Sum of Distances in Tree

문제 유형 : dfs, 재귀, 트리, 다이나믹 프로그래밍

문제 난이도 : Hard

 

문제

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.

 

방향이 없는 트리가 0~n-1까지 라벨이 붙은 n개의 노드와 n-1개의 간선으로 구성됩니다.

당신은 정수 n을 받습니다.

edges[i] = [a_i, b_i]는 a_i와 b_i가 연결되어 있음을 말합니다.

answer[i] = i번노드에서 다른 모든 노드들의 거리들의 합을 말합니다.

조건을 만족하는 answer배열을 반환하시오.

 

풀이

처음에는 dp[src][dst]로 두고 일반적인 재귀로 풀려고 했다.

그러나, 이 문제의 핵심은 '트리'이며, 간선이 n-1개 라는 것이다.

 

트리를 기준으로 생각해서, 루트에서 각 노드들까지의 거리의 합은 결국 자식 서브트리들의 크기와 관련이 있다.

0을 트리의 루트라고 하고, dfs를 돌려보자.

자식노드들의 서브트리의 크기를 재귀적으로 누적하면 된다.

 

이제, 0에서 각 노드들의 합은 구해졌다.

이제 나머지 자식노드들의 답을 구하면 된다.

dfs를 한번 더 해보자.

각 서브트리의 자식노드들의 답은, 서브트리의 루트의 모든 합 - 자식 노드의 서브트리 크기 + n - 자식노드의 서브트리 크기이다.

(<자식 노드의 서브트리 크기>를 서브트리의 루트의 모든 합에서 뺴는 이유는, 서브트리에서 루트보다 한 칸 더 가까운 노드들이기 때문이다.

(n - <자식노드의 서브트리 크기>)는 루트 노드로부터 한 칸씩 더 멀기 때문이다.)

 

코드

C++

class Solution {
private:
    vector<vector<int>> graph;
    vector<int> subtreeSize;
    vector<int> ans;
    int n;

    void dfs(int node, int parent) {
        subtreeSize[node] = 1;
        for (int child : graph[node]) {
            if (child != parent) {
                dfs(child, node);
                subtreeSize[node] += subtreeSize[child];
                ans[node] += ans[child] + subtreeSize[child];
            }
        }
    }

    void dfs2(int node, int parent) {
        for (int child : graph[node]) {
            if (child != parent) {
                ans[child] = ans[node] - subtreeSize[child] + n - subtreeSize[child];
                dfs2(child, node);
            }
        }
    }

public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        this->n = n;
        graph.resize(n);
        subtreeSize.resize(n, 0);
        ans.resize(n, 0);

        for (const auto& edge : edges) {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }

        dfs(0, -1);
        dfs2(0, -1);

        return ans;
    }
};