Count 4-cycles [CSAcademy #65 B]

https://csacademy.com/contest/round-65/task/count-4-cycles/

2つのN頂点の木がある。
それぞれの木の同じ頂点間に辺を追加で張り、2つの木を合成する。
この無向グラフの中で長さ4のサイクルの個数を答えよ。

解法

長さ4のサイクルの数を数えるということだが、もうちょっとわかりやすく制約を言い換えてみる。
「頂点a,bに対して2つの木それぞれに辺がある場合に合成すると長さ4のサイクルになる」
となるので、今回の問題は上の条件をみたす(a,b)の組の数が答えになる。
これは片方の辺について全探索をして、もう片方にその辺があるかを確認する。
setなどでもう片方の辺を保持しておいて、チェックしてもいいし、下の実装のように隣接リストの中身をソートしておいてlower_boundで探してもいい。
後者の方が早そうだが、setの方が実装が簡単なのでオススメ。

int N;
vector<int> E1[101010], E2[101010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N - 1) {
        int a, b; cin >> a >> b;
        a--; b--;
        E1[a].push_back(b);
        E1[b].push_back(a);
    }
    rep(i, 0, N - 1) {
        int a, b; cin >> a >> b;
        a--; b--;
        E2[a].push_back(b);
        E2[b].push_back(a);
    }
    rep(i, 0, N) sort(all(E1[i]));
    rep(i, 0, N) sort(all(E2[i]));

    int ans = 0;
    rep(i, 0, N) {
        fore(j, E1[i]) if (i < j) {
            auto ite = lower_bound(all(E2[i]), j);
            if (ite != E2[i].end() and *ite == j) ans++;
        }
    }
    cout << ans << endl;
}