kmjp's blog

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

101 Hack 47 : E. Summing in a Tree

問題の意図を読み間違えて無駄に苦戦。
https://www.hackerrank.com/contests/101hack47/challenges/summing-in-a-tree

問題

N頂点の根付き木が与えられる。
頂点vのレベルとは、根頂点からの距離を意味する。
f(l,k) := 頂点vのうち、サブツリー内にレベルlの頂点がk個以上あるものの個数
を意味する。

木の高さ(レベルの最大値)をHとする。
数列A[i]に対し、 \sum_{i=0}^H f(i,A_i)を求めよ。

解法

A[i]=0となるレベルiについては、すべての頂点が条件を満たすことが明確であるため以後除外する。

いわゆる頂点のマージテクを使い、以下のデータを葉から根に挙げていこう。

  • S(v) := 頂点vのサブツリーにおいて、レベルlの頂点がA[l]個以上あるようなlのset
  • M(v) := 頂点vのサブツリーにおいて、レベルlの頂点がA[l]個以下あるようなlとその登場回数のmap

子頂点から親頂点にS(v),M(v)をどんどんマージしていく過程で、後者でレベルlの頂点がA[l]個を超えた場合、そのlはM(v)からS(v)に移動する。
各頂点vは|S(v)|だけ答えに寄与することになる。

HackerRankでは5*10^5回再帰したらSegmentation Faultを起こしたため、以下は無理やり再帰関数を回避している。

int N,H;
int D[505050];
vector<int> E[505050];
int A[505050];
ll ret;

set<int> S[505050];
map<int,int> M[505050];
int dif[505050];
vector<int> cand[505050];

void dfs(int cur) {
	if(A[D[cur]]==1) {
		S[cur].insert(D[cur]);
	}
	else if(A[D[cur]]>1) {
		M[cur][D[cur]]++;
	}
	
	FORR(e,E[cur]) {
		if(S[cur].size()<S[e].size()) swap(S[cur],S[e]);
		FORR(r,S[e]) S[cur].insert(r);
		S[e].clear();
		if(M[cur].size()<M[e].size()) swap(M[cur],M[e]);
		FORR(r,M[e]) {
			if(S[cur].count(r.first)) continue;
			M[cur][r.first] += r.second;
			if(M[cur][r.first]>=A[r.first]) {
				S[cur].insert(r.first);
				M[cur].erase(r.first);
			}
		}
		M[e].clear();
	}
	
	ret += S[cur].size();
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>H;
	FOR(i,N-1) {
		cin>>x;
		E[x].push_back(i+1);
	}
	FOR(i,H+1) {
		cin>>A[i];
		if(A[i]==0) ret += N;
	}
	
	queue<int> Q;
	Q.push(0);
	while(Q.size()) {
		x=Q.front();
		Q.pop();
		FORR(r,E[x]) {
			D[r]=D[x]+1;
			Q.push(r);
		}
	}
	
	FOR(i,N) cand[D[i]].push_back(i);
	
	for(i=H;i>=0;i--) FORR(e,cand[i]) dfs(e);
	
	cout<<ret<<endl;
	
}

まとめ

このぐらいの再帰関数でSEGV起こすの勘弁。