Bonsai Grafting [「みんなのプロコン 2019」決勝 B]

https://atcoder.jp/contests/yahoo-procon2019-final/tasks/yahoo_procon2019_final_b

前提知識

• 全方位木DP（これじゃなくてもいいけど）

解説

https://atcoder.jp/contests/yahoo-procon2019-final/submissions/4358833

まず、それぞれの木について、ある頂点から最遠の頂点までの距離を計算する。
これは、全方位木DPで達成できる。
全方位木DPを理解していれば、この構築は難しくないので、全方位木DPを知らない場合は、
そちらをまずは学習しよう。
（自分は本戦で、木の中心のライブラリとLCAで頂点間の距離を求めるアレをして盛大にバグった）
　
これを使って、問題を解いていく。
N頂点の木のある頂点cuに注目して考えてみる。
頂点cuのこの木における最遠の距離をdist1[cu]としておこう。
頂点cuはM頂点と結ぶ可能性があり、結んだ先を頂点toとすると、
最も長い距離は
1. 両方の木の直径の大きい方（変数maと置く）
2. 新しくできた直径（dist1[cu] + 1 + dist2[to]）
のどちらかになる。
M通りある結び方のうち、場合1になる組合せ、場合2になる組合せを高速に求めたい。
このためにBITを使う。
bit[i] := distがiである頂点数
これを使うと、場合1になるときはdist1[cu] + 1 + dist2[to]≦maのときであり、
dist2[to]≦ma - dist1[cu] - 1と変形できる。
つまり、ng=bit[0...ma - dist1[cu] - 1]が場合1になる組合せで、M-ngが場合2である。
場合1は後で処理するとして、otherに追加しておこう。
場合2はこの連結によってdist1[cu]がM-ng回答えに追加されるので、これを答えに足そう。

1. 1の部分は最後にまとめて足して、+dist2[to]の部分はM頂点の木でも同様のことを行ったときに足される。

これをM頂点の木の方でも行う。
　
otherはN側とM側で2回ダブってカウントされているので、実際はother/2通り場合1になっている。
よって、ma*other/2を足す。
最後に、+1の部分が残っているので、other/2通りではない場合に+1が発生するので、NM-other/2を足すと答え。

void _main() {
    int N, M;
 
    Reroot rr1, rr2;
 
    N = rr1.Load();
    vector<ll> dist1 = rr1.build();
 
    M = rr2.Load();
    vector<ll> dist2 = rr2.build();
 
    ll ma = -1;
    rep(i, 0, N) chmax(ma, dist1[i]);
    rep(i, 0, M) chmax(ma, dist2[i]);
 
    BIT<ll> bit1(101010), bit2(101010);
    rep(i, 0, N) bit1.add(dist1[i], 1);
    rep(i, 0, M) bit2.add(dist2[i], 1);
    ll other = 0;
 
    ll ans = 0;
    rep(i, 0, N) {
        ll ng = 0;
        if(0 <= ma - dist1[i]) ng = bit2.get(0, ma - dist1[i]);
        other += ng;
        ans += dist1[i] * (M - ng);
    }
    rep(i, 0, M) {
        ll ng = 0;
        if (0 <= ma - dist2[i]) ng = bit1.get(0, ma - dist2[i]);
        other += ng;
        ans += dist2[i] * (N - ng);
    }
    ans += ma * other / 2;
    ans += 1LL * (1LL * N * M - other / 2);
    cout << ans << endl;
}