AtCoder Beginner Contest 173, F : Intervals on Tree

問題文

　https://atcoder.jp/contests/abc173/tasks/abc173_f

問題概要

　$N$ 頂点の木があって，頂点は $1$ から $N$ の整数で番号付けされている．また，$i$ 番目の辺は 2 点 $u_i, v_i$ を結んでいる．
　整数 $L, R$ ($1 \leq L \leq R \leq N$) について，関数 $f$ を以下のように定義する．

• $f( L, R ) :=$ 頂点集合 $S = [ L, R ]$ について，誘導部分グラフ $G[ S ]$ の連結成分の個数

　この値を全ての区間について足し上げた値，$$\sum_{ L = 1 }^{ N } \sum_{ R = L }^{ N } f( L, R )$$ を求めよ．

制約

• $1 \leq N \leq 2 \times 10^5$

解法

　まず簡略化したバージョンを考えてみようということで，誘導部分グラフをとるのではなく，頂点集合だけをとって辺を全く使わない場合を考えてみます．この場合は，連結成分の個数というのは頂点数と一致するので，すなわち区間の幅と同じということになります．幅 $1$ の区間は $[ 1, 1 ], [ 2, 2 ], \dots [ N, N ]$ と $N$ 個あり，幅 $2$ の区間は $[ 1, 2 ], [ 2, 3 ], \dots, [ N - 1, N ]$ と $N - 1$ 個あり……，と続くので，全体では $$\sum_{ i = 1 }^{ N } i ( N + 1 - i )$$ 個となります*1．
　次に，辺について考えてみます．ある一つの辺について，以下のようなことが言えます．

• その辺を含む誘導部分グラフの個数分，連結成分の個数の総和が減る

これは，その辺を使うと両端点が接続され，かつ，元のグラフが木であるので既に連結であるということはない，ということから言えます．従って，各辺についてその辺を含む誘導部分グラフの個数を求め，先程の $\sum$ の値から引いていけば答えが求まります．この問題での誘導部分グラフのとり方は区間がベースになっているので，辺 $( u, v )$ （一般性を失わずに $u < v$ とする）を含む誘導部分グラフの個数というのは，区間 $[ u, v ]$ を左右に広げる方法の総数に一致します．これは，$u$ 側を左に伸ばす方法と $v$ 側を右に伸ばす方法の組合せなので，$u \times ( N + 1 - v )$ 通りとなります．
　結局，グラフを構築したりはせずに辺の端点の番号を使って算数（？）をすると答えが求まります．計算量は $\Theta( N )$ です．

コード

int main()
{
	IN( int, N );
	LL res = 0;
	REP( i, 1, N + 1 )
	{
		res += LL( i ) * ( N + 1 - i );
	}

	REP( N - 1 )
	{
		IN( int, u, v );
		if ( v < u )
		{
			swap( u, v );
		}
		--u;

		const int l = u + 1;
		const int r = N - v + 1;
		res -= LL( l ) * r;
	}
	cout << res << endl;

	return 0;
}

*1:この $\sum$ をもっと簡単に書けるとかそういう話は割愛．Wolfram 先生に訊くと答えだけ教えてもらえるけどわたしには導出できなかった……（ア）