kmjp's blog

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

AtCoder ABC #283 (ユニークビジョンプログラミングコンテスト2022 冬) : Ex - Popcount Sum

久々に来たな…。
https://atcoder.jp/contests/abc283/tasks/abc283_h

問題

整数N,M,Rが与えられる。
N以下の正整数のうち、Mで割ってR余るものについて、popcountの総和を求めよ。

解法

テストケース数が多いため、以下では間に合わなかった。

  • Mが大きい場合、R+k*M≦Nとなる整数を総当たり
  • Mがそこそこ小さい場合、桁DP

ここではfloor_sumを使う。
XをN以下の正整数のうち、Mで割ってR余るものの個数(floor((N-R)/M)+1)とする。

Mで割ってR余る整数において、2^iのbitが立っている数を数え上げることを考えよう。
これができれば、iを総当たりすればpopcountの総和が得られる。

floor_sum(X,pow(2,i),M,R)は、N以下の正整数のうち、Mで割ってR余るものについて、さらに2^iで割った値の総和を得られる。
求めたいのは2^iで割った値を2で割った値の総和だが、それはfloor_sum(X,pow(2,i),M,R) - 2*floor_sum(X,pow(2,i+1),M,R)とすればよい。

int T;
ll N,M,R;
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>>T;
	while(T--) {
		cin>>N>>M>>R;
		
		ll ret=0;
		
		ll pre=floor_sum((N-R+M)/M,1,M,R);
		FOR(i,60) {
			ll a=floor_sum((N-R+M)/M,2LL<<i,M,R);
			ret+=pre-a*2;
			pre=a;
		}
		
		cout<<ret<<endl;
	}
}

まとめ

初手桁DPに走って失敗した…。
微妙にMが多くてダメだった。