TopCoder SRM 775 : Div1 Easy Div2 Hard IterateOverACube

なんとかレート増をキープ。
https://community.topcoder.com/stat?c=problem_statement&pm=15895

問題

(0,0,0)～(N-1,N-1,N-1)のN^3種類のTupleを以下の順で並べる。K番目は何か。

• 総和が小さいものを前に持ってくる
• 総和が等しい場合、辞書順に並べる

解法

タプルの総和を0～(3N-3)まで総当たりし、それぞれの総和を満たすタプル数を総和がKに達するまで数えていこう。
これは包除原理の応用で1つの総和に対し組み合わせはO(1)で求められる。
総和が確定したら、同様に第1要素を総当たりしつつ第2・3要素の組み合わせをO(1)で数えていけばよい。

ll comb(int P_,int Q_) {
	if(P_<0 || Q_<0 || Q_>P_) return 0;
	ll p=1,q=1;
	Q_=min(Q_,P_-Q_);
	for(int i=1;i<=Q_;i++) p=p*P_, q=q*i,P_--;
	
	return p/q;
}
ll hcomb(int P_,int Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}


ll pat[3030303];

class IterateOverACube {
	public:
	vector <int> findCell(int N, long long index) {
		int i;
		ll ret=0;
		index++;
		FOR(i,3*N-2) {
			pat[i]=comb(i+2,2);
			if(i>=N) pat[i]-=3*comb(i+2-N,2);
			if(i>=2*N) pat[i]+=3*comb(i+2-2*N,2);
			ret+=pat[i];
			
			if(index<=pat[i]) {
				for(int x=0;x<=min(i,N-1);x++) {
					int left=i-x;
					if(left<=N-1) {
						int num=left+1;
						if(index<=num) return {x,index-1,left-(index-1)};
						index-=num;
					}
					else {
						if(left>2*N-2) continue;
						int num=2*N-1-left;
						if(index<=num) {
							for(int y=0;y<=N-1;y++) {
								int z=i-x-y;
								if(z>=0 && z<=N-1) {
									index--;
									if(index==0) return {x,y,z};
								}
							}
						}
						index-=num;
						
					}
					
				}
			}
			else {
				index-=pat[i];
			}
			
		}
		
		return {};
	}
}

まとめ

問題設定はシンプルだけど、割と神経使う問題。