問題
自分で考えても分からなかったですが、典型問題なので解説ACしました。
解法
全ての頂点同士の最短経路を求めようとすると、TLEになります。
そこで、それぞれの辺を通る頂点の組み合わせは何通りか？を考えます。
上の赤い辺の場合
(1, 2), (1, 3), (1, 4), (5, 2), (5, 3), (5, 4）
の6通り。
この組み合わせを高速に計算することを考えると、上のように辺の手前と向かいでグループを分けると良いことに気付きます。
それぞれのグループから要素を一つずつ選ぶ場合の数が、辺を通る頂点の組み合わせになります。
上の例では、3 * 2 = 6通り
DFSの実装
木を探索するときは深さ優先探索(DFS)を使うと良いことが多いです。
今回の場合では、それぞれの要素が子要素をいくつ持つかを求めたいので、帰りがけに子要素の数を回収するように実装します。
// Nは木の要素数
void dfs(int pos, int pre, vector<long long> &dp, vector<vector<long long>> &tree, long long &ans, long long N) {
dp[pos] = 1; // 自分
for (auto n : tree[pos]) {
if (n == pre) continue; // 逆流防止
dfs(n, pos, dp, tree, ans, N);
// ここより下は帰りがけに通る。
dp[pos] += dp[n]; // 子要素の数を回収
}
ans += dp[pos] * (N - dp[pos]); // (自分 + 子要素) * それ以外の要素
}