問題の意図を読み間違えて無駄に苦戦。
https://www.hackerrank.com/contests/101hack47/challenges/summing-in-a-tree
問題
N頂点の根付き木が与えられる。
頂点vのレベルとは、根頂点からの距離を意味する。
f(l,k) := 頂点vのうち、サブツリー内にレベルlの頂点がk個以上あるものの個数
を意味する。
木の高さ(レベルの最大値)をHとする。
数列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起こすの勘弁。