【AtCoder】경프전형 039 Tree Distance

문제

스스로 생각해도 몰랐습니다만, 전형적인 문제이므로 해설 AC했습니다.

해법

모든 정점끼리의 최단 경로를 구하려고 하면 TLE가 됩니다.
거기서, 각각의 변을 통과하는 정점의 조합은 몇 가지인가? 생각합니다.


위의 빨간 변의 경우
(1, 2), (1, 3), (1, 4), (5, 2), (5, 3), (5, 4)
6 거리.



이 조합을 빠르게 계산하는 것을 고려하면 위와 같이 변 앞과 맞은편으로 그룹을 나누면 좋은 것을 알 수 있습니다.
각각의 그룹으로부터 요소를 하나씩 선택하는 경우의 수가, 변을 통과하는 정점의 조합이 됩니다.
위의 예에서 3*2=6

DFS 구현

나무를 탐색할 때는 깊이 우선 탐색(DFS)을 사용하면 좋은 것이 많습니다.
이 경우에는 각 요소가 자식 요소를 몇 개 가지고 있는지 묻고 싶기 때문에 돌아가서 자식 요소의 수를 회수하도록 구현합니다.

// Nは木の要素数
void dfs(int pos, int pre, vector<long long> &dp, vector<vector<long long>> &tree, long long &ans, long long N) {
    dp[pos] = 1; // 自分

    for (auto n : tree[pos]) {
        if (n == pre) continue; // 逆流防止

        dfs(n, pos, dp, tree, ans, N);
        // ここより下は帰りがけに通る。

        dp[pos] += dp[n]; // 子要素の数を回収
    }

    ans += dp[pos] * (N - dp[pos]); // (自分 + 子要素) * それ以外の要素
}

경쟁 프로 전형적인 90 질문에 대해

AtCoder에서의 실력 업을 목표로하자!　～경전전형 90문～