floor sum と聞いて！！

問題へのリンク

問題概要

$1$ 以上 $N$ 以下の整数のうち、 $M$ で割って $R$ 余るものを考える。そのような整数の popcount 値の総和を求めよ。

( $T$ ケース与えられる)

制約

$1 \le T \le 10^{5}$
$1 \le M \le N \le 10^{9}$
$0 \le R \lt M$

考えたこと

TL を見て、floor sum と知った上での考察。

「popcount の総和を求めよ」という設定自体から、とりあえず各桁ごとに求めればよさそうという推測は立つ。

$k$ 桁目の値が $1$ になるものの個数を求めるために、その条件を言い換えていくことにする。まずよくあるのは、

$x$ の $k$ 桁目が $1$ である
$\Leftrightarrow$ $x$ を $2^{k+1}$ で割ったあまりが $2^{k}$ 以上 $2^{k+1}$ 未満である

という言い換え。この言い換えは、たとえば次の問題で本質的に活躍する。この言い換えは苦手で、少し前に頑張って意識したので、今回はスムーズに出た。

drken1215.hatenablog.com

これは、次のように言い換えられる。

$\displaystyle \bigl\lfloor \frac{x + 2^{k}}{2^{k+1}} \bigr\rfloor - \bigl\lfloor \frac{x}{2^{k+1}} \bigr\rfloor = 1$

つまり、元の問題は次のように言い換えられる。

$0$ 以上 $N$ 以下の $M$ で割って $R$ あまる整数 $x$ について、

$\displaystyle \bigl\lfloor \frac{x + 2^{k}}{2^{k+1}} \bigr\rfloor - \bigl\lfloor \frac{x}{2^{k+1}} \bigr\rfloor$

の総和を求めよ。

さらに

$x = Mt + R$ とおくと、 $t$ の条件は

$\displaystyle 0 \le t \le \lfloor \frac{N-R}{M} \rfloor$

となる。最終的な答えは、 $n = \lfloor \frac{N-R}{M} \rfloor + 1$ とおいて、

$\displaystyle \sum_{t = 0}^{n-1} \bigl\lfloor \frac{Mt + R + 2^{k}}{2^{k+1}} \bigr\rfloor - \sum_{t = 0}^{n-1} \bigl\lfloor \frac{Mt + R}{2^{k+1}} \bigr\rfloor$

と表せる。この値は floor sum で求められる。計算量は $O((\log N)^{2})$ となる。

コード

#include <bits/stdc++.h>
using namespace std;

// sum_{i=0}^{n-1} floor((a * i + b) / m)
// O(log(n + m + a + b))
long long floor_sum(long long n, long long a, long long b, long long m) {
    if (n == 0) return 0;
    long long res = 0;
    if (a >= m) {
        res += n * (n - 1) * (a / m) / 2;
        a %= m;
    }
    if (b >= m) {
        res += n * (b / m);
        b %= m;
    }
    if (a == 0) return res;
    long long ymax = (a * n + b) / m, xmax = ymax * m - b;
    if (ymax == 0) return res;
    res += (n - (xmax + a - 1) / a) * ymax;
    res += floor_sum(ymax, m, (a - xmax % a) % a, a);
    return res;
}

int main() {
    int CASE;
    cin >> CASE;
    while (CASE--) {
        long long N, M, R;
        cin >> N >> M >> R;
        long long range = (N - R) / M + 1;

        long long res = 0;
        for (int k = 0; k <= 30; ++k) {
            long long upper = floor_sum(range, M, R + (1LL<<k), 1LL<<(k+1));
            long long lower = floor_sum(range, M, R, 1LL<<(k+1));
            res += upper - lower;
        }
        cout << res << endl;
    }
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 283 Ex - Popcount Sum (橙色, 600 点)

問題概要

制約

考えたこと

さらに

コード