TopCoder SRM 763 : Div2 Hard ETSums

いやこの入力形式止めて…。
https://community.topcoder.com/stat?c=problem_statement&pm=15563&rd=17614

問題

N頂点の根付き木が与えられる。
各頂点vには以下のパラメータが設定される。
C(v) := コスト
T(v) := ID。根頂点を1とし、そこからDFSの訪問順にインクリメントされる。

ET(v)を以下のように定義する。根頂点からvへのパスを、s0,s1,...skとすると

• ET(v) = sum(C(si)*i**T(v))

ET(v)の和を求めよ。

解法

C(v)がどのぐらい寄与するかを考える。
vの根からの距離をdとすると、vのSubTree内の頂点集合をSとするとき、

• 

さて、T(c)が任意の値を取るとするとこれは普通に求めることはできず、O(N^2)かかりかねない。
しかし、ここでT(c)の部分は、連番の範囲しかとらないことに注意しよう。
すると、総和の部分は実は単なる等比数列の和となり容易に求めることができる。

vector<ll> C;
vector<int> E[202020];
int T[202020];
ll mo=1000000007;
int id;
ll ret;

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

class ETSums {
	public:
	void dfs(int cur,int d) {
		int L=id++;
		FORR(e,E[cur]) dfs(e,d+1);
		int R=id;
		
		if(d==1) {
			(ret+=C[cur]*(R-L))%=mo;
		}
		else if(d>1) {
			ll a=modpow(d,L)*(modpow(d,R-L)+mo-1)%mo*modpow(d-1)%mo;
			(ret+=C[cur]*a)%=mo;
		}
		
		
	}
	int findSumOfETSums(int N, vector <int> parent, vector <int> cost, int D, int seed, int MX) {
		vector<ll> A(2*N);
		A[0]=seed;
		int i;
		for(i=1;i<2*N;i++) A[i] = (A[i-1] * 1103515245LL + 12345) %2147483648;
		for(i=parent.size();i<N;i++) parent.push_back(A[i]%min(i,D)+i-min(i,D));
		for(i=cost.size();i<N;i++) cost.push_back(A[N+i]%MX);
		C.clear();
		FOR(i,N) C.push_back(cost[i]);
		FOR(i,N) E[i].clear();
		for(i=1;i<N;i++) E[parent[i]].push_back(i);
		id=1;
		ret=0;
		dfs(0,0);
		return ret;
	}
}

まとめ

連番になるのに気が付くまで割と手間取った。