yukicoder : No.3432 popcount & sum (Hard)

これは普通に自力で解けた。
https://yukicoder.me/problems/no/3432

問題

f(a,b)を、aとbのpopcountが等しいときは(a and b)、等しくないときは0とする。
整数Nが与えられるので、を998244353で割った余りを求めよ。

解法

g(d,p) := N以下の正整数のうち、2^dのbitが立っていて、かつpopcountがpとなるものの数
とすると、d,pを総当たりしながら(2^d)*g(d,p)*(g(d,p)+1)/2を解に計上すればよい。

g(d,p)は桁DPの要領で求めることができる。

ll N;
const ll mo=998244353;

ll dp[61][61][2][61];

ll dfs(int d,int must,int le,int lef) {
	if(lef<0) return 0;
	if(d<0) return lef==0;
	if(dp[d][must][le][lef]>=0) return dp[d][must][le][lef];
	int i;
	ll ret=0;
	FOR(i,2) {
		if(le==0&&i==1&&(N&(1LL<<d))==0) continue;
		if(must==d&&i==0) continue;
		ret+=dfs(d-1,must,le||(i==0&&(N&(1LL<<d))),lef-i);
	}
	return dp[d][must][le][lef]=ret%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	ll ret=0;
	MINUS(dp);
	
	FOR(i,60) {
		FOR(j,61) {
			ll num=dfs(60,i,0,j)%mo;
			(ret+=(num*(num+1)/2%mo)*((1LL<<i)%mo))%=mo;
		}
	}
	cout<<ret<<endl;
}

まとめ

シンプルな問題設定だけど、意外に解いた記憶なかったな。