D - Masked Popcount Editorial by MMNMM


\(\displaystyle f(L,R)=\sum _ {i = L} ^ R\operatorname{popcount}(i\mathop{\&}M)\) の値を求めやすいような整数の組 \((L,R)\) を用いてこの問題を解くことを考えます。

\((L,R)=(2 ^ ed,2^e(d+1)-1)\) とすると、\[\operatorname{popcount}((L+i)\mathop{\&}M)+\operatorname{popcount}((R-i)\mathop{\&}M)\ (0\leq i\lt2 ^ e)\] が \(i\) によらず一定であることがわかります。 よって、このとき \[f(L,R)=(\operatorname{popcount}(L\mathop{\&}M)+\operatorname{popcount}(R\mathop{\&}M))\times2 ^ {e-1}\] が成り立ちます。

あとは、このような形の \(f(L,R)\) を用いて \(f(0,N)\) を求めればよいです。 これは、ABC355 E - Guess the Sum と同じ形式ですが、最小回数である必要はないのでたかだか \(2\lceil\log _ 2N\rceil\) 個の区間の和として表しても十分高速です。

実装例は以下のようになります。

#include <iostream>
#include <bit>
#include <atcoder/modint>

int main() {
    using namespace std;
    using modint = atcoder::static_modint<998244353>;
    unsigned long N, M;
    cin >> N >> M;
    ++N;
    modint ans{};
    while(N){
        const auto i{N & -N};
        ans += modint::raw(popcount(M & N - i) + popcount(M & N - 1)) * i;
        N ^= i;
    }
    cout << (ans / 2).val() << endl;
    return 0;
}

posted:
last update: