AtCoder ABC #147 : F - Sum Difference

いつもより正答者が少ないね。
https://atcoder.jp/contests/abc147/tasks/abc147_f

問題

整数N,X,Dが与えられる。
初項X、公差Dの等差数列の先頭N要素を考える。
それらを2人で任意に分け合い、その総和の差を考えると、何通り考えられるか。

解法

D=0の場合はコーナーケースとして先に片づけておく。X=0かどうかは注意しよう。
D＜0の場合は等差数列の先頭N要素の並びを反転させ、以下D＞0とする。

Nが小さいので、一人目が何要素を取るかを総当たりしよう。
仮にk要素取るとする。
一人目の値の総和は、kX + {(0～(N-1)D)までのうちk要素の和}の形で表現される。
この時後ろの項は、最小値から最大値まですべての値をDごとに取ることができる。
(これは容易に証明できる。最大値でないなら、どこか取る要素よりD大きな取ってない要素が残っているはずなので、そこを交換すればよい)

差を考えると、2Dごとにある区間の値をすべて取ることができるということになる。
よってk=0～Nまで総当たりすると、そのような区間が(N+1)通りできる。
なのでmod 2Dで分類して和集合を取ればよい。

ll N;
ll X,D;

map<ll,vector<pair<ll,ll>>> M;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>X>>D;
	
	if(D==0) {
		if(X==0) cout<<1<<endl;
		else cout<<(N+1)<<endl;
		return;
	}
	
	if(D<0) {
		X+=(N-1)*D;
		D=-D;
	}
	
	ll tot=D*N*(N-1)/2;
	FOR(i,N+1) {
		ll ma=D*N*i-D*i*(i+1)/2;
		ll mi=D*i*(i-1)/2;
		
		ll lef=X*i-X*(N-i)+mi-(tot-mi);
		ll ri=X*i-X*(N-i)+ma-(tot-ma)+2*D;
		ll cat=(lef%(2*D)+2*D)%(2*D);
		assert((ri-lef)%(2*D)==0);
		M[cat].push_back({lef,ri});
	}
	
	ll ret=0;
	FORR(m,M) {
		vector<pair<ll,ll>> V=m.second;
		sort(ALL(V));
		ll L=-1LL<<62,R=L;
		FORR(v,V) {
			if(v.first<=R) {
				R=max(R,v.second);
			}
			else {
				ret+=(R-L)/(2*D);
				L=v.first;
				R=v.second;
			}
		}
		ret+=(R-L)/(2*D);
	}
	cout<<ret<<endl;
	
}

まとめ

シンプルな問題設定で良かった。