yukicoder : No.2206 Popcount Sum 2

気付いてしまえば楽なのだが。
https://yukicoder.me/problems/no/2206

問題

以下の問題がTケース与えられる。

整数N,Mが与えられるので、以下に答えよ。
2^N未満の正整数のうち、popcountがM以下であるものの総和を求めよ。

解法

(0-originで)2進数表記でd桁目が1になるものの組み合わせは、C(N-1,0)+C(N-1,1)+....+C(N-1,M)通りである。
それぞれ2^dだけ解に寄与すると考えると、(1+2+4+...+2^(N-1))*(C(N-1,0)+C(N-1,1)+....+C(N-1,M))が解となる。
前者は等比数列なので簡単に計算できるとして、後者は前計算ありでもO(M)かかる。

f(n,m) = C(n,0)+C(n,1)+...+(n,m)
とする。f(n,m)からf(n+1,m),f(n-1,m),f(n,m+1),f(n,m-1)への遷移はO(1)でできる。
よって、T個のテストケースをMo's Algorithmの要領で並べ替え、f(n,m)を順次求めて行こう。

int T;
int N[202020],M[202020];
const ll mo=998244353;
ll ret[202020];

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	vector<vector<int>> V;
	cin>>T;
	FOR(i,T) {
		cin>>N[i]>>M[i];
		V.push_back({N[i]/500,M[i],i});
	}
	int CN=0,CM=0;
	ll cur=1;
	sort(ALL(V));
	FORR(v,V) {
		i=v[2];
		int TN=N[i]-1,TM=M[i]-1;
		while(TN>CN) cur=(cur*2+mo-comb(CN++,CM))%mo;
		while(TN<CN) cur=(cur+comb(--CN,CM))*(mo+1)/2%mo;
		while(TM>CM) (cur+=comb(CN,++CM))%=mo;
		while(TM<CM) (cur+=mo-comb(CN,CM--))%=mo;
		ret[i]=cur%mo;
	}
	
	
	FOR(i,T) {
		ll sum=ret[i]*(modpow(2,N[i])-1)%mo;
		cout<<sum<<endl;
	}
}

まとめ

Moに考えが至らず…。