競プロ典型90題 - 039 Tree Distance（★5）

問題　https://atcoder.jp/contests/typical90/tasks/typical90_am
解説　kyopro_educational_90/039.jpg at main · E869120/kyopro_educational_90 · GitHub

【問題】
重みのない木が与えられる。
任意の2頂点の最短パス長の総和を求めよ。

【典型ポイント】
主客転倒
辺をカットして頂点を２グループ化

【方針】
辺に注目し、各辺をカットして頂点を２グループに分けた時の頂点数の積が、その辺の「答えへの貢献度」
貢献度の和が答えになる。
辺をカットしたときのそれぞれの頂点数は、[子ノードと、その子孫のノード数]をAとすると、A * (N - A)となる。

【感想】
知らんかった。主客転倒って何

int main(){
    int N; cin >> N;
    vector<vector<int>> G(N);
    rep(i, N-1){
        int a, b; cin >> a >> b;
        G[--a].push_back(--b);
        G[b].push_back(a);
    }

    vector<ll> vc(N); //各頂点の、その頂点を含む子孫の数
    auto dfs1 = [&](auto&& self, int v, int p) -> ll{
        vc[v] = 1;
        for(auto e:G[v]){
            if(e==p)continue;
            vc[v] += self(self, e, v);
        }
        return vc[v];
    };
    dfs1(dfs1, 0, -1);

    auto dfs2 = [&](auto&& self, int v, int p) -> ll{
        ll res = 0;
        for(auto e:G[v]){
            if(e==p)continue;
            res += self(self, e, v) + vc[e] * (N - vc[e]);
        }
        return res;
    };
    cout << dfs2(dfs2, 0, -1) << ln;
}