Codeforces #759 : Div2 E. Frequency Queries

結構手間取ってるな。
https://codeforces.com/contest/1591/problem/E

問題

根付き木が与えられる。
また、各頂点には整数値が設定されている。

以下の各クエリに答えよ。
クエリは頂点vと、2値L,Kで与えられる。
vから根までの最短パス上の頂点に設定された整数値からなる数列を考える。
ここから登場頻度がL以上のものについて、頻度順に並べたときに、少ない方からK番目のものを答えよ。

解法

DFSで根付き木を探索しながら、BITなどで頻度がいくつのものが何個あるかを保持しよう。
このBITを二分探索すれば、解に相当する値の頻度がわかる。

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	void set(int e,V v) { add(e,v-val[e]);}
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
};
BIT<int,21> bt;



int T,N,Q;
int A[1010101];
int C[1010101];
int P[1010101];
vector<int> E[1010101],ev[1010101];
unordered_set<int> S[1010101];

int V[1010101],L[1010101],K[1010101];
int ret[1010101];

void dfs(int cur) {
	bt.add(C[A[cur]],-1);
	S[C[A[cur]]].erase(A[cur]);
	
	C[A[cur]]++;
	
	S[C[A[cur]]].insert(A[cur]);
	bt.add(C[A[cur]],1);
	
	FORR(e,ev[cur]) {
		ll sum=bt(L[e]-1);
		if(sum+K[e]>bt(N)) {
			ret[e]=-1;
		}
		else {
			int x=bt.lower_bound(sum+K[e]);
			ret[e]=*S[x].begin();
		}
		
	}
	FORR(e,E[cur]) dfs(e);
	
	bt.add(C[A[cur]],-1);
	S[C[A[cur]]].erase(A[cur]);
	
	C[A[cur]]--;
	
	S[C[A[cur]]].insert(A[cur]);
	bt.add(C[A[cur]],1);
	
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	FOR(i,1010000) {
		S[0].insert(i);
	}
	bt.add(0,1010000);
	
	cin>>T;
	while(T--) {
		cin>>N>>Q;
		FOR(i,N) {
			cin>>A[i];
			E[i].clear();
			ev[i].clear();
		}
		for(i=1;i<N;i++) {
			cin>>x;
			E[x-1].push_back(i);
		}
		FOR(i,Q) {
			cin>>V[i]>>L[i]>>K[i];
			ev[V[i]-1].push_back(i);
		}
		dfs(0);
		
		
		FOR(i,Q) cout<<ret[i]<<" ";
		cout<<endl;
	}
}

まとめ

入出力10^6個オーダーはちょっと重いな…。