TopCoderOpen 2020 East Asia : Hard CannonballPyramids

参加してたらHard落としてた…。
https://community.topcoder.com/stat?c=problem_statement&pm=16449&rd=18266

問題

S番目のピラミッド数とは、ボールをピラミッド状にS段重ねたときのボース総数(1^2+2^2+...+S^2)とする。
ここで正整数Nが与えられる。
Nを6個以下のピラミッド数の和で表せ。

解法

N以下のピラミッド数はO(N^(1/3))個あるので、3・4個集めればN通りの数はだいたい表せそうである。
そこであと2個で微調整しよう。

まず最初の2個は0～2番目までのピラミッド数を総当たりする。
残り4つは、半分全列挙しよう。ピラミッド数は1500個ぐらいなので、2個ずつ総当たりしても間に合う。
ある数を4つの数の和で表す際に半分全列挙を使うのは蟻本の最初の方に出てくるよね。

class CannonballPyramids {
	public:
	vector <int> build(int B) {
		vector<pair<int,ll>> P;
		ll sum=0;
		int i;
		for(i=0;i<=10000;i++) {
			sum+=i*i;
			P.push_back({i,sum});
			if(sum>B) break;
		}
		
		vector<vector<int>> cand={
			{0,0,0},
			{0,1,1},
			{0,2,5},
			{1,1,2},
			{1,2,6},
			{2,2,10},
		};
		
		map<int,pair<int,int>> C;
		FORR(p,P) FORR(q,P) {
			ll s=p.second+q.second;
			C[s]={p.first,q.first};
			
			FORR(c,cand) {
				ll lef=B-(s+c[2]);
				if(C.count(lef)) {
					vector<int> ret;
					ret.push_back(p.first);
					ret.push_back(q.first);
					ret.push_back(C[lef].first);
					ret.push_back(C[lef].second);
					ret.push_back(c[0]);
					ret.push_back(c[1]);
					
					sort(ALL(ret));
					while(ret[0]==0) ret.erase(ret.begin());
					return ret;
					
				}
				
			}
			
			
		}
		
		return {};
		
	}
}

まとめ

Nの上限が無限に大きくできるとして、どこぐらいまでなら6個で表せるんだろうな。