AtCoder ABC #313 (第四回日本最強プログラマー学生選手権-予選-) : G - Redistribution of Piles

FよりもだいぶAC数多め。
https://atcoder.jp/contests/abc313/tasks/abc313_g

問題

N個の皿があり、i番目の皿にはA[i]個の石が乗っている。
また、空の袋がある。

以下の処理を任意回数行えるとする。

• 石が1個以上乗っている皿から、1個ずつ石を回収して袋に入れる。
• 袋にN個以上石がある場合、N個取り出して各皿に1個ずつ配る

最終的にあり得る石の数の列は何通りか。

解法

後者の処理のあと前者の処理をすると互いに打ち消しあうだけで意味がない。
よって、前者の処理を何回か行った後、後者を行うことを考える。
後者を行えるのは、前者の処理を行った後、floor(袋の石の数/N)+1である。

よって前者の処理回数毎にfloorの部分の和を求めればよい。
これはfloor_sumで求めることができる。
前者の処理回数毎に、袋に入れられる石の数がだんだん減っていくので、O(N)回floor_sumを呼ぼう。

int N;
ll A[202020];
const ll mo=998244353;

ll floor_sum(ll N,ll M,ll A,ll B) {
	// sum(i=0...N-1) floor((A*i+B)/M)
	
	ll ret=0;
	if(B>=M) ret+=N*(B/M), B%=M;
	if(A>=M) ret+=N*(N-1)/2*(A/M), A%=M;
	
	ll Y=(A*N+B)/M;
	if(Y==0) return ret;
	//floor(Y/M)に達するX
	ll X=Y*M-B;
	//Xの右側はY個ずつ
	ret+=(N-(X+A-1)/A)*Y;
	// 90度回転、Y=Nのラインは無視する
	ret+=floor_sum(Y,A,M,(A-X%A)%A);
	return ret;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	sort(A,A+N);
	ll ret=(A[0]+1)%mo;
	ll sum=1LL*A[0]*N;
	for(i=1;i<N;i++) if(A[i]>A[i-1]) {
		ret+=floor_sum(A[i]-A[i-1],N,(N-i),sum+(N-i)+N);
		ret%=mo;
		
		sum+=1LL*(A[i]-A[i-1])*(N-i);
	}
	cout<<ret<<endl;
}

まとめ

海外コンでfloor_sum使う例見たことないけど、どこかあるのかな。