mod Nで等差数列/等比数列の和を計算

最近知ったちょいテクをメモっておく。

等差数列の和

等差数列の和をmod Nで計算することを考える。
初項a, 等差d, 項数nとすれば、求める和は、

sum = n (2a + (n-1) d) / 2

である。この値はmod Nではどうなるか？
gcd(2, N) = 1であれば、分子に2の逆元をかければいい。

Nが偶数の場合はどうすればいいか？一般に q | pのとき

p / q (mod N) = p (mod N * q) / q

が成り立つ。

よって、(分子 mod 2N) / 2を計算すればよい（※）。

ll sum(ll a, ll d, ll n, ll mod) {
    mod *= 2;
    ll ret = (2 * a + d * (n-1)) % mod;
    ret = ret * n % mod;
    return ret / 2;
}

※実は、n、2a + (n-1) dのいずれかは必ず2で割り切れるので、予め割れる方を割っておけば赤字のテクニックを知らなくても計算できる。

等比数列の和

初項a, 等比r, 項数nとする。求める和は、

sum = a(r^n - 1) / (r - 1)

である。法Nの値によっては、(r-1)の逆元が存在しない場合もある。
この場合も、赤字のテクニックを使うことで、和を計算できる。

または、n項の和をn/2項ずつの和に分割して計算することも出来る。

ll pow(ll x, ll p, ll mod) {
    ll ret = 1;
    while (p) {
        if (p & 1)
            ret = ret * x % mod;
        x = x * x % mod;
        p >>= 1;
    }
    return ret;
}

ll sum(ll a, ll r, ll n, ll mod) {
    if (n == 1)
        return a % mod;

    ll x = sum(a, r, n/2, mod);
    ll ret = (x + pow(r, n/2, mod) * x) % mod;
    if (n & 1)
        ret = (a + r * ret) % mod;
    return ret;
}