Codeforces #302 Div1 D. Road Improvement

0除算にやられた。
http://codeforces.com/contest/543/problem/D

問題

1～N番の町が道路で連結されており、その形状は木を成すグラフをなっている。

初期状態ではすべての辺は状態が悪い。
町の１つvが首都であるとすると、各町から首都に至る経路において、通る必要がある状態が悪い道路の数を１個以下になるように幾つかの道路を舗装したい。

1～N番の各町が首都vである場合それぞれについて、舗装する道路の組み合わせ数を求めよ。

解法

グラフを根付木と見なしてDFSを２回行う。
１回目ではsubtreeにおける条件を満たす組み合わせを求めていき、２回目で親側の状態を伝搬させる。

頂点xにおいて、subtree上の頂点からxに至る経路が題意を満たすような組み合わせをdp(x)とする。
dp(x)のうち、1つは全道路を舗装するケースであり、残りは1箇所以上状態が悪い道路が残っているケースである。
よって、頂点xの子頂点をyとすると

• x-yの間を舗装しない場合、yのsubtreeからxに至る経路が条件を満たすのは1通り(y以下に状態の悪い道路が1つもない場合)
• x-yの間を舗装する場合場合、yのsubtreeからxに至る経路が条件を満たすのはdp(y)通り。

よってで計算できる。

次に頂点xの親側の経路に置いて同様に条件を満たす組み合わせをdp2(x)とする。
すると問題の解はdp(v)*dp2(v)となる。

頂点xの親頂点をzとすると、dp2(x) = dp(z)*dp2(z) / (1+dp(x))となるので２回目のDPでこれを求めていけば良い。
なお、この式には除算が登場する。
一見(1+dp(x))のmod 10^9+7における逆元を取ればよさそうだが、1+dp(x) が0になるケースで失敗する。
よってzに置いてzの子頂点wに対し(1+dp(w))の累積積やSegTreeを用いて(1+dp(x))だけを除いた積を高速に計算できるようにしておかなければならない。

int N;
vector<int> E[202000];
ll mo=1000000007;

struct RemCumProd {
	vector<ll> le,ri,v;
	void init(vector<ll> V) { v=V; init(); }
	void add(ll d) { v.push_back(d); }
	void init() {
		le.clear(); ri.clear();
		le.push_back(1),ri.push_back(1);
		for(int i=0;i<v.size();i++) le.push_back(le.back()*v[i]%mo), ri.push_back(ri.back()*v[v.size()-1-i]%mo);
	}
	ll operator[](int x) { // if x<0 || x>=N, return all prod
		if(x<0 || x>=le.size()-1) return le.back();
		return le[x]*ri[le.size()-2-x]%mo;
	}
};

ll ppat[202000];
RemCumProd patp[202000];

void dfs(int cur) {
	FORR(r,E[cur]) {
		dfs(r);
		patp[cur].add(1+patp[r][-1]);
	}
	patp[cur].init();
}

void dfs2(int cur, ll p) {
	ppat[cur]=p;
	int i;
	FOR(i,E[cur].size()) dfs2(E[cur][i],1+p*patp[cur][i]%mo);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	cin>>N;
	FOR(i,N-1) cin>>x, E[x-1].push_back(i+1);
	dfs(0);
	dfs2(0,1);
	
	FOR(i,N) _P("%I64d ",ppat[i]*patp[i][-1]%mo);
	_P("\n");
}

まとめ

今回の木DPで反省したので、累積積から1個だけ除いたものを返すコードをライブラリ化した。