kmjp's blog

競技プログラミング参加記です

yukicoder : No.2834 Work to Play

式の立て方次第で難しさが変わりそう。
https://yukicoder.me/problems/no/2834

問題

整数N,A,Kが与えられる。
座標1~Nのそれぞれにおいて、座標iにはA*i個のカートがある。

カートをK個まで詰める台車を用いて以下の手順でカートを原点に回収したい。

  • まず、K個以上カートがある場所があれば、そこまで行ってカートをK個積んで原点に戻る。
  • どの場所もK個未満になったら、原点から一番遠いカートのある場所まで行き、そこからK個まで遠い順に道中カートを台車に積みつつ戻る。

総移動距離を求めよ。

解法

両方のステップをそれぞれ求める。
あらかじめA,KはGCDで割っておき、A,Kは互いに素な状態にしておく。

1つ目のステップについて。
座標Kx+yについて、y=0~(K-1)それぞれxに対し移動量・往復回数が線形に増えるので、それぞれxの1乗和・2乗和を用いて計算しよう。

2つ目のステップは、座標2K毎に回収の仕方を繰り返すことになるので、2K分だけ愚直にシミュレートしよう。

ll N,A,K;
const ll mo=998244353;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>A>>K;
	ll g=__gcd(A,K);
	A/=g;
	K/=g;
	
	ll ret=0;
	FOR(j,K) {
		__int128 r=(N-j)/K;
		if(j>N) continue;
		r%=mo;
		ret+=K*(A%mo)%mo*((r*(r+1)*(2*r+1)/6)%mo);
		ret+=(A%mo)*j%mo*((r*(r+1)/2)%mo);
		ret+=((__int128)A*j/K*K%mo)*((r*(r+1)/2)%mo);
		ret+=((__int128)A*j/K%mo)*j%mo*(r+1)%mo;
		ret%=mo;
	}
	
	ll cur=0;
	FOR(i,2*K) {
		ll a=__int128(N-i)*A%K;
		if(a>cur) {
			a-=cur;
			cur=K-a;
			
			ll t=N-i;
			ll s=(N-i)%(2*K);
			(ret+=((__int128)(s+t)*((t-s)/(2*K)+1)/2)%mo)%=mo;
			
		}
		else {
			cur-=a;
		}
	}
	
	ret=ret*2%mo;
	cout<<ret<<endl;
	
}

まとめ

自力だとこう綺麗にすんなり式立てられなさそう…。