Codeforces Round #513 by Barcelona Bootcamp (rated, Div. 1 + Div. 2) E. Sergey and Subway

Problem - E - Codeforces

問題概要

Nノードの無向木が与えられる．辺 (p,q) および辺 (q,r) が存在するようなp,q,rについて，pとrの間に辺を張る．

あらゆるノードのペア間の最短距離の総和を求める． ただし辺のコストは全て1とする．

解法概要

辺を追加する前の木と後の木では，ノード間の最短距離は次のように変化する．

• 辺を追加する前に最短距離が 2d (dは正の整数) だったノードのペアは，辺を追加した後の最短距離がdになる．
• 辺を追加する前に最短距離が 2d+1 だったノードのペアは，辺を追加した後の最短距離がd+1になる．

したがって，辺を追加する前の最短距離の総和に，最短距離が2d+1だったノードのペアの数を足すと，あとは2で割るだけで辺を追加した後の最短距離の総和が求まる．

なお，辺を追加する前の木での最短距離の総和は，各辺(u,v)を切って木を2つに分けた場合を考え (u側の木を構成するノードの数) × v側の木を構成するノードの数) を求め，すべての辺についてその値の総和をとって最後に2で割れば良い．(2で割らないとノードのペア (u,v) と (v,u) を別々のものとして考えることになる．)

最短距離が2d+1だったノードのペア数は，辺の両端が同じ色にならないように各ノードを白と黒に塗り分けていくテクで計算できそう．

下記のソースコードでは途中までペア中のノードの順番を気にしない ((u,v) と (v,u) を別々のものとして考える) 方針で計算し，最後にまとめて2で割っている．

ソースコード

#include<iostream>
#include<vector>

using namespace std;

typedef long long ll;

const int MAX_N = 200000;

int n;
ll ans = 0LL;
vector<int> edges[MAX_N];
ll dp[MAX_N][2];
bool used[MAX_N];

void dfs(int node, bool is_black){
    used[node] = true;

    for(int child: edges[node]){
        if(used[child]) continue;
        dfs(child, true^is_black);
        ll n_children = dp[child][0] + dp[child][1];
        ans += n_children * (n-n_children); // node<->childを通る回数 (ペアの順番を気にしない)
        dp[node][0] += dp[child][0];
        dp[node][1] += dp[child][1];
    }
    dp[node][is_black]++;
    return;
}

int main(){
    std::ios::sync_with_stdio(0); cin.tie(0);
    cin >> n;
    for(int i=0;i<n-1;i++){
        int from, to;
        cin >> from >> to;
        from--; to--;
        edges[from].push_back(to);
        edges[to].push_back(from);
    }

    dfs(0, 0);
    ans += dp[0][0] * dp[0][1]; // 距離が奇数のペアの数 (ペアの順番を気にしない)

    cout << ans/2 << endl; // ペアの順番を気にしないで計算してきたので2で割る
}

所感

木DP 自力で解けた ことがない (DP自体 そもそも苦手)

参考

NarenJitter氏のコード