TopCoder SRM 683 Div2 Hard SubtreesCounting

これ既に世の中のどこかで出てそうだ。
https://community.topcoder.com/stat?c=problem_statement&pm=14179

問題

木を成すグラフが与えられる。
このグラフのサブグラフとは、このグラフのうち１つ以上の連結頂点群とそれらを結ぶ辺からなるグラフの事を意味する。

与えれたグラフの全サブグラフにおける頂点数の総和を求めよ。

解法

グラフを根付き木と見なし、各頂点を根とするサブグラフの組み合わせの総和を求めよう。
それには以下の値を木DPで更新していけばよい。S(v)の総和が解となる。
A(v) := 頂点vを根とするサブグラフの組み合わせ数
S(v) := 頂点vを根とするサブグラフの頂点数の和

一つも子頂点を持たない頂点の場合、A(v)=S(v)=1は明確である。
頂点vといくつかの子頂点を処理した状態のA(v),S(v)が判明しているとする。
次の子頂点c以下をvを根とするサブグラフに連結する場合、A(v),S(v)がどうなるかを考える。

vにcを連結しない場合を考慮すると、A(v) = A(v) * (1+A(c))で更新できる。
問題はS(v)。

vを頂点とするA(v)個のサブグラフの頂点数配列を例えばV(v)とし、cを頂点とするA(c)個のサブグラフの頂点数配列を例えばV(c)とする。
求めるS(v)は以下の和である。

• cと連結しない場合、(V(v)[0]+V(v)[1]+...V(v)[A(v)-1])
• cのサブグラフのうちA(c)個中(0-indexで)0番目と連結する場合、(V(v)[0]+V(c)[0])+(V(v)[1]+V(c)[0])+...(V(v)[A(v)-1]+V(c)[0]) = (V(v)[0]+V(v)[1]+...V(v)[A(v)-1]) + A(v)*V(c)[0]
• cのサブグラフのうちA(c)個中(0-indexで)1番目と連結する場合、(V(v)[0]+V(c)[1])+(V(v)[1]+V(c)[1])+...(V(v)[A(v)-1]+V(c)[1]) = (V(v)[0]+V(v)[1]+...V(v)[A(v)-1]) + A(v)*V(c)[1]
• …

つまり、S(v) = S(v)*(A(c)+1) + A(v)*S(c)で更新できる。

ll A[101010];
vector<int> E[101010];
ll pat[101010];
ll tot[101010];

ll ret;
ll mo=1000000007;

class SubtreesCounting {
	public:
	void dfs(int cur,int pre) {
		pat[cur]=tot[cur]=1;
		FORR(r,E[cur]) if(r!=pre) {
			dfs(r,cur);
			(tot[cur] = tot[cur]*(pat[r]+1)+pat[cur]*tot[r])%=mo;
			(pat[cur]*=(1+pat[r]))%=mo;
		}
		ret+=tot[cur];
	}
	
	int sumOfSizes(int n, int a0, int b, int c, int m) {
		int i,j;
		A[0]=a0;
		for(i=1;i<=n-2;i++) A[i]=(b*A[i-1]+c)%m;
		
		FOR(i,n) E[i].clear();
		for(i=1;i<n;i++) {
			j=A[i-1]%i;
			E[i].push_back(j);
			E[j].push_back(i);
		}
		ret=0;
		dfs(0,-1);
		return ret%=mo;
		
	}
}

まとめ

Div1 Mediumと入力数の制限が10^5だったのでCodeforcesっぽいな。