けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 378 F - Add One Edge 2 (1D, 水色, 500 点)

evima さんの別解で解いた。

問題概要

頂点数  N の単純な木が与えられる。この木に辺を 1 本追加して得られるグラフのうち、次の条件を満たすものの個数を求めよ。

  • 単純グラフである
  • サイクルをちょうど 1 つ含み、そのサイクルに含まれる頂点の次数がすべて 3 である

制約

  •  3 \le N \le 2 \times 10^{5}

考えたこと

木の各頂点の次数を求めていこう。このとき、次数が 3 であるような頂点を連結成分分解しよう。

このとき、条件を満たすグラフは次のように言い換えられるのだ。


同一の連結成分に対して、その連結成分に接続している頂点のうち、次数が 2 であるものを 2 つ選んで辺を張る


このように言い換えられれば、解法は自然に導かれる。

  • 木の頂点のうち、次数が 3 である部分の連結成分を抽出していく
  • 各連結成分について、それに接続している次数 2 の頂点数を  K として、 {}_{K}\mathrm{C}_{2} を足していく

この解法の計算量は  O(N) となる。

コード

#include <bits/stdc++.h>
using namespace std;

using Graph = vector<vector<int>>;
int main() {
    int N, u, v;
    cin >> N;
    Graph G(N);
    for (int i = 0; i < N-1; i++) {
        cin >> u >> v, u--, v--;
        G[u].push_back(v), G[v].push_back(u);
    }

    vector<bool> seen(N, false);
    vector<int> comp;
    auto rec = [&](auto rec, int v) -> void {
        seen[v] = true;
        if (G[v].size() == 3) comp.push_back(v);
        for (auto v2 : G[v]) {
            if (seen[v2] || G[v2].size() != 3) continue;
            rec(rec, v2);
        }
    };

    long long res = 0;
    for (int i = 0; i < N; i++) {
        if (seen[i] || G[i].size() != 3) continue;
        comp.clear();
        rec(rec, i);

        long long num = 0;
        for (auto v : comp) {
            for (auto v2 : G[v]) {
                if (G[v2].size() == 2) num++;
            }
        }
        res += num * (num - 1) / 2;
    }
    cout << res << endl;
}