AtCoder ARC #108 : F - Paint Tree

方針を微妙に間違えてた。
https://atcoder.jp/contests/arc108/tasks/arc108_f

問題

N頂点の木を成す無向グラフが与えられる。
各点を白黒で塗り分けたとする。
グラフのスコアを、同色の頂点の最長距離とする。
塗り分け方2^N通りにおけるスコアの総和を求めよ。

解法

まず木の直径を求めよう。
直径の両端が同色の場合、最長距離は直径で確定する。
以下、両端の色が異なるケースを考えよう。

f(d) := 同色の頂点の最長距離がdである塗り分け方の個数
とすると、解はd*f(d)である。f(d)を直接求めるのは難しいので、
g(d) := 同色の頂点の最長距離がd以下である塗り分け方の個数
とするとf(d)=g(d)-g(d-1)で求めることができる。

以降、g(d)を数えることを考える。
直径の両端以外の頂点において、両端からの距離が(x,y)だったとする。

• max(x,y)≦dなら、この点は白黒どちらでもg(d)に数え上げる条件を満たす。
• min(x,y)＞dなら、この点は白黒どちらでもg(d)に数え上げる条件を満たさない。
• min(x,y)≦d＜max(x,y)なら、片方の色でのみg(d)に数え上げる条件を満たす。

とすると、各頂点はdに応じて0,1,2通りのいずれかの色の選び方が考えられる。
あとは累積和など使って各頂点における選び方の積を取ればg(d)を求めることができる。

int N;
vector<int> E[402020];
const ll mo=1000000007;
ll p2[505050];
vector<int> D[2];

ll pow2[202020];
ll pow0[202020];
ll dp[202020];

pair<int,int> farthest(int cur,int pre,int d,vector<int>& D) {
	D[cur]=d;
	pair<int,int> r={d,cur};
	FORR(e,E[cur]) if(e!=pre) r=max(r, farthest(e,cur,d+1,D));
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	p2[0]=1;
	FOR(i,404040) p2[i+1]=p2[i]*2%mo;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	D[0].resize(N);
	D[1].resize(N);
	auto v1=farthest(0,0,0,D[0]);
	auto v2=farthest(v1.second,v1.second,0,D[0]);
	v1=farthest(v2.second,v2.second,0,D[1]);
	x=v1.second;
	y=v2.second;
	
	FOR(i,N) {
		if(i==x||i==y) continue;
		int a=D[0][i];
		int b=D[1][i];
		if(a>b) swap(a,b);
		pow2[b]++;
		pow0[a-1]++;
	}
	
	for(i=N;i>=0;i--) pow0[i]+=pow0[i+1];
	FOR(i,N+1) {
		pow2[i+1]+=pow2[i];
		if(pow0[i]==0) dp[i]=p2[pow2[i]];
	}
	ll ret=0;
	for(i=N;i>=1;i--) {
		(dp[i]+=mo-dp[i-1])%=mo;
		(ret+=i*dp[i])%=mo;
	}
	ret=ret*2%mo;
	(ret+=p2[N-1]*v1.first)%=mo;
	
	cout<<ret<<endl;
	
}

まとめ

方針を誤ったときのリカバリ方法を考えないとな。