kmjp's blog

競技プログラミング参加記です

Codeforces #426 Div1 D. Red-Black Cobweb

Codeforces

こういう式変形もあるのね。
http://codeforces.com/contest/833/problem/D

問題

N頂点の木を成す無向グラフが与えられる。
各辺は赤黒の色が振られ、またパラメータxを持つ。

この木におけるパスのうち、赤黒辺の多い方が少ない方の２倍以下に収まっているものを考える。
そのようなすべてのパスにおけるパラメータxの積を答えよ。

解法

重心分解を行いつつ処理していく。
さて、赤黒辺の数に関する条件を考えてみる。
重心からある頂点vまでの辺で赤辺・黒辺が(R,B)だったとする。
重心から別の向きにある別の頂点v'に対し同様に赤黒辺の数が(R',B')だったとする。

重心を挟みパスv-v'が答えにカウントされるのは、2*min(R+R',B+B') ≧ max(R+R',B+B') である。
この条件を言い換えてみよう。
まず赤辺が十分多くなければならないので2*(R+R')≧B+B'とすると2*R-B ≧ B'-2*R'となる。
よって事前に訪問済み頂点v'に対しB'-2*R'および対応する重心からのxの積を計算し、SegTreeで管理しておけば、vをDFSしながら対応するv'の数とxの積を求めることができる。

2*(R+R')≧B+B'だけだと、黒辺が少なすぎるケースを排除できてないので、そこから2*B-R＜R'-2*B'のケースを除けばよい。

以下のコードでは以下4つのBITを使った。

  • indexを-(2*R-B)とし、対応する頂点の数を保持するBIT
  • indexを-(2*R-B)とし、重心から対応する頂点までのxの積の積を保持するBIT
  • indexを-(2*B-R)とするもので上２つと同様のもの。
int N;
vector<pair<int,int>> EV[101010];
ll X[101010],C[101010];
int vis[101010];
int NV[101010];
ll mo=1000000007;
ll ret=1,retr=1;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {V s=0;e+=(1<<(ME-2));while(e) s+=bit[e-1],e-=e&-e; return s;}
	V add(int e,V v) { e+=(1<<(ME-2)); while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
template<class V, int ME> class BITm {
public:
	V bit[1<<ME];
	BITm(){ for(int i=0;i<1<<ME;i++) bit[i]=1;}
	V operator()(int e) {V s=1;e+=(1<<(ME-2));while(e) (s*=bit[e-1])%=mo,e-=e&-e; return s;}
	V add(int e,V v) { e+=(1<<(ME-2)); while(e<=1<<ME) (bit[e-1]*=v)%=mo,e+=e&-e;}
};

BIT<int,20> R2,B2;
BITm<ll,20> R2m,B2m;
vector<pair<pair<int,int>,ll>> evs,evt;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll rev(ll a) {
	ll r=1, n=mo-2;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

int dfs(int cur,int pre) {
	NV[cur]=1;
	FORR(e,EV[cur]) if(e.first!=pre && vis[e.first]==0) NV[cur]+=dfs(e.first,cur);
	return NV[cur];
}

int dfs2(int cur,int pre,int TNV) {
	int tot=1;
	int ok=1;
	FORR(e,EV[cur]) if(e.first!=pre && vis[e.first]==0) {
		int c = dfs2(e.first,cur,TNV);
		if(c!=-1) return c;
		tot += NV[e.first];
		if(NV[e.first]*2>TNV) ok=0;
	}
	if((TNV-tot)*2>TNV) ok=0;
	if(ok) return cur;
	return -1;
}

void dfs3(int cur,int pre,int R,int B,ll V) {
	evt.push_back({{R,B},V});
	
	int x=2*R-B;
	
	(ret *= modpow(V,R2(x)))%=mo;
	(ret *= R2m(x))%=mo;
	
	x = 2*B-R;
	(retr *= modpow(V,B2(220000)-B2(x)))%=mo;
	(retr *= B2m(220000)*rev(B2m(x))%mo)%=mo;
	
	FORR(e,EV[cur]) if(e.first!=pre && vis[e.first]==0) dfs3(e.first,cur,R+(C[e.second]==0),B+(C[e.second]==1),V*X[e.second]%mo);
}

void split(int cur,int id) {
	int TNV = dfs(cur,-1);
	if(TNV==0) return;
	int center=dfs2(cur,-1,TNV);
	vis[center]=1;
	evs.push_back({{0,0},1});
	R2.add(0,1);
	B2.add(0,1);
	
	FORR(e,EV[center]) if(vis[e.first]==0) {
		dfs3(e.first,center,C[e.second]==0,C[e.second]==1,X[e.second]);
		
		FORR(r,evt) {
			evs.push_back(r);
			R2.add(-(r.first.first*2-r.first.second),1);
			R2m.add(-(r.first.first*2-r.first.second),r.second);
			B2.add(-(r.first.second*2-r.first.first),1);
			B2m.add(-(r.first.second*2-r.first.first),r.second);
		}
		evt.clear();
	}
	FORR(r,evs) {
		R2.add(-(r.first.first*2-r.first.second),-1);
		B2.add(-(r.first.second*2-r.first.first),-1);
		R2m.add(-(r.first.first*2-r.first.second),rev(r.second));
		B2m.add(-(r.first.second*2-r.first.first),rev(r.second));
	}
	evs.clear();
	FORR(e,EV[center]) if(vis[e.first]==0) split(e.first,id+1);
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y>>X[i]>>C[i];
		EV[x-1].push_back({y-1,i});
		EV[y-1].push_back({x-1,i});
	}
	split(0,1);
	
	cout<<ret*rev(retr)%mo<<endl;
	
}

まとめ

本番時間あっても重心分解は思いついたとして2*R-B ≧ B'-2*R'あたりの変形ができずダメそう。