NoiminのNoise

競技プログラミング (多め) とWeb (たまに) ,自然言語処理 (ブログではまだ)。数式の書き方を一気に KaTeX に変えようとして記事を全削除してインポートし直すなどしたので,過去にブックマークされた記事は URL が変わってしまっている可能性があります…….

AtCoder Beginner Contest 173 F - Intervals on Tree

問題: F - Intervals on Tree

公式解説: 解説 pdf

問題概要

$N (\leq 2 \times 10^5)$ 頂点の無向木がある。整数 $1 \leq L \leq R \leq N$ について,関数 $f(L, R)$ を次のように定義する。

  • $S$ を$L$ 以上 $R$ 以下の番号のついた頂点からなる集合とする。頂点集合 $S$ と両端が $S$ に属する辺からなる部分グラフの連結成分の数を $f(L, R)$ とする。

$\sum_{L=1}^{N} \sum_{R=L}^N f(L, R)$ を求めよ。

解法概要

コンテスト中は「木を根付き木と考えて連結成分の値となりうる回数について全方位木 DP とか……?」と考えて泥沼に陥ってしまった (これでも解けないわけではない? わからん) が,もっとシンプルに考えられる。

両端が $S$ に属する辺を1本使うと,頂点集合 $S$ を使った部分グラフの連結成分の数は必ず 1 つ減る。つまり,辺を使わない場合の連結成分 (= 頂点) 数の合計から,各辺が使われることによって減る連結成分の数の合計を引けば良い。

辺が1本もない場合を考えてみる。辺がない場合の連結成分の個数は頂点数に等しいので

$\sum_{L=1}^{N} \sum_{R=L}^N (R-L+1)$

$= \sum_{L=1}^{N} \sum_{r=1}^{N+1-L} r$

($m = N + 1 - L$ として)

$= \sum_{L=1}^{N} \sum_{r=1}^{m} r = \sum_{L=1}^{N} \frac{m(m+1)}{2} = \frac{1}{2} ( \frac{1}{6}N(N+1)(2N+1) + \frac{N(N+1)}{2} )$

($L = 1$ のとき $m = N + 1 - 1 = N$, $L = N$ のとき $m = N + 1 - N = 1$).

ここから,各辺によって連結成分が減らされる回数の合計を求めたい。これは,例えば辺 $(u, v)$ があって $u \lt v$ の場合,この辺によって連結成分が減らされる回数は $u (n+1-v)$ 回である。これは, $u$ が含まれるような $L$ の選び方は $u$ 通り,$v$ が含まれるような $R$ の選び方は $n+1-v$ 通りあるので,$u$ と $v$ が両方含まれる $S$ の選び方の通り数が $u (n+1-v)$ 通りとなるためである。これをすべての辺について計算して,先ほど求めた辺がない場合の連結成分の個数の合計から引けば良い。

ソースコード

#include <iostream>
#include <algorithm>

using namespace std;

using ll =  long long;

int main() {
    ll n, u, v;
    cin >> n;
    ll ans = (n * (n+1) * (2*n+1) / 6 + n * (n+1) / 2) / 2;
    for(int i=0;i<n-1;++i) {
        cin >> u >> v;
        ans -= min(u, v) * (n - max(u, v) + 1);
    }
    cout << ans << endl;
}

所感

発想次第でとてもシンプルになる問題。オーバーフローにだけ注意。