Sum of Maximum Weights [AtCoder Beginner Contest 214 D]

https://atcoder.jp/contests/abc214/tasks/abc214_d

前提知識

• （主客転倒）
• UnionFind

解説

https://atcoder.jp/contests/abc214/submissions/25065159

見た目はかなり難しい。
方針が分かれば実装はかなり簡単で済むのだが、考え方は難しい。

何がモチベーションとなるか

今回の問題はすべての頂点間に対する処理になるので、すべての頂点間の列挙だけでO(N²)となる。
愚直にやっていくと、

すべての頂点間の組み合わせについて f(i,j) の総和

が答えになる。
だが、ここで役立つ条件としてf(i,j)は重みの種類だけパターンがある。
つまり、f(i,j)を全探索することは可能であるということである。
よって、

すべてのf(i,j)について f(i,j)×それを満たす頂点間の組み合わせ の総和

を計算することにしよう。
この組み合わせを高速に計算できれば、計算量が間に合いそうだ。
このように観点を変えて計算することを主客転倒と（一部の人は）呼んでいる。
まあ、名前にそれほど意味は無くて汎用的に使える思考法という感じである。
あの〜、お詫びと言っては何ですけどちょっと数え上げでよく見るらしい「主客転倒」の解説今から書くんで… - physics0523's 精進ログ

その組み合わせをいかに求めるか

さて、この組み合わせをいかに求めるかという部分に主軸を置いて考えてみよう。
f(i,j)がとある辺kの重みw[k]である場合について考えてみる。
u[k] <- w[k] -> v[k]
という辺があって、これを単純に通る頂点の組について考えると、
u[k]側の頂点の個数×v[k]側の頂点の個数となる。
だが、これでは別のw[k]以上の頂点が含まれる可能性がある。
よって、求めたい組合せは

u[k]側のw[k]以下の辺を使って到達可能な頂点の個数×v[k]側のw[k]以下の辺を使って到達可能な頂点の個数

となる。この式まで理解することが完全理解につながる。
よくわからない場合はグラフを書いて理解してみよう。

到達可能な頂点の個数

これだけを見て考えてみると、一般的にはUnionFindによって実現ができそうだなと思い浮かぶ。
その方向性で考えると、工夫をすることでUnionFindを使ってこの個数を高速に求めることができる。

UnionFindをうまく使って要求を満たす

w[k]の辺を処理する場合は、w[k]以下の辺について結合処理がなされている必要がある。
よって、辺の評価は重みが小さいものから順番に、
式の計算→結合処理→式の計算→結合処理→…
という風にやっていけば、結合部分（連結成分）に含まれる頂点数は高速に求まるので、
要求されている計算を行うことができる。

int N;
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    vector<tuple<int, int, int>> edges;
    rep (i, 0, N - 1) {
        int u, v, w;
        cin >> u >> v >> w;
        u--; v--;
        edges.push_back({ w, u, v });
    }
    sort(all(edges));

    UnionFind uf(N);
    ll ans = 0;
    fore(edge, edges) {
        int u, v, w;
        tie(w, u, v) = edge;

        ans += 1LL * w * uf.getValues(u) * uf.getValues(v);
        uf(u, v);
    }
    cout << ans << endl;
}