【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 問～