yukicoder : No.942 プレゼント配り

解けはしたけどだいぶ手間取った。
https://yukicoder.me/problems/no/942

問題

整数N,Kが与えられる。NはKの倍数である。
1～Nの数字をN/K個ずつK個のグループに分ける。
総和が一緒になるようにできるか。

解法

K=NまたはK=1の時は自明なので、以下1＜K＜Nとする。
まず各グループの総和を等しくできるか、1～Nの和がKの倍数かを判定しよう。

C=N/Kとする。
Cが偶数なら、最初のK個は昇順、次のK個は降順…と交互に割り当てればよい。
そうでない場合、K＜NよりCは3以上の奇数である。
そこで最初の3K個だけうまく均等に割り当てれば、残り(C-3)個ずつの割り当ては偶数と同じ手順で行ける。

3K個の配分については、後ろK個を降順に並べ、先頭K個を1個飛ばしでおいて、最後に真ん中を均等になるように配置すればよい。
K=5の時こんな感じ。

15 1 8
14 3 7
13 5 6
12 2 10
11 4 9

ll N,K;

vector<ll> R[202020];

void ok() {
	int i;
	cout<<"Yes"<<endl;
	FOR(i,K) {
		FORR(c,R[i]) cout<<c<<" ";
		cout<<endl;
	}
	exit(0);
	return;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	ll sum=N*(N+1)/2;
	if(sum%K) return _P("No\n");
	if(K==N && K!=1) return _P("No\n");
	
	int P=N/K;
	
	if(P%2==0 || P==1) {
		FOR(x,P) {
			FOR(y,K) {
				if(x%2==0) R[y].push_back(x*K+y+1);
				else R[K-1-y].push_back(x*K+y+1);
			}
		}
		ok();
	}
	if(P>=3) {
		ll X=3*K*(3*K+1)/2/K;
		
		FOR(y,K) R[y].push_back(3*K-y);
		FOR(y,K) {
			if(y<=K/2) R[y].push_back(1+y*2);
			else R[y].push_back((y-K/2)*2);
			R[y].push_back(X-R[y][0]-R[y][1]);
		}
		
		for(x=3;x<P;x++) {
			FOR(y,K) {
				if(x%2==0) R[y].push_back(x*K+y+1);
				else R[K-1-y].push_back(x*K+y+1);
			}
		}
		ok();
	}
	
	
	cout<<"No"<<endl;
}

まとめ

Cが奇数の場合が鬼門。