公式解説と同じような方針で \(O(T\log(N))\) で解くことができます。公式解説にあった前計算は \(\Theta(\log^2(N))\) から \(\Theta(\log(N))\) に改善されます。

\(0\) 以上 \(N\) 未満の整数の集合を \(2\) 冪のサイズの集合に綺麗に分けます。

具体的には、\(N = 22 = 2^{4} + 2^{2} + 2^{1}\) のとき、\([0, 16), [16, 20), [20, 22)\) のように分けます。

これらの集合に含まれる整数は2進数で以下のように表現できます。

• \([0, 16)\)

0????

• \([16, 20)\)

100??

• \([20, 22)\)

1010?

それぞれの集合に対して、popcount が \(K\) であるような整数の総和を求めれば良いです。

\([L, L + 2^{a})\) に対する答えを求めます。(\(L\) は \(2^{a + 1}\) で割り切れます。)

\(\mathrm{popcount}(L) = c\) とすると、そのような整数の数は \(\binom{a}{K - c}\) です。これは \(a\) 個の ? からちょうど \(K - c\) 個が 1 であるような振り分け方と同じだからです。

また、popcount が \(K\) であるような \(\binom{a}{K - c}\) 個の整数からランダムに数を選んだとき、その期待値は \(L + (2^{a} - 1) \cdot \frac{K - c}{a}\) です。なぜなら、それぞれの ? について、1 になる確率が独立に \(\frac{K - c}{a}\) だからです。

よって、\([L, L + 2^{a})\) に対する答えは \(\binom{a}{K - c}\left(L + (2 ^ {a} - 1)\cdot \frac{K - c}{a}\right)\) になります。

適切な前計算のもとで、それぞれの区間に対して答えを \(O(1)\) で求められるので、テストケースごとに \(O(\log(N))\) で求められます。\(\mathrm{popcount}(N) = K\) のときは答えに \(N\) を加算する必要があることに注意してください。

c++ による実装例

ABC 過去問の類題