等比数列の総和 (mod) (Math/mod_sum_of_geometric_sequence.hpp)

View this file on GitHub
View document part on GitHub
Last update: 2024-09-30 17:31:16+09:00
Include: #include "Math/mod_sum_of_geometric_sequence.hpp"

sum_of_geometric_sequence

mint sum_of_geometric_sequence(mint a, mint r, long long n)

初項$a$、公比$r$、項数$n$からなる等比数列の部分和 $\sum_{i=0}^{n-1} a \cdot r^i$ を計算し、その値を返します。

制約

$0 < n$

計算量

$O(\log n)$

Verified with

test/Math/mod_sum_of_geometric_sequence/abc293-e.test.cpp

Code

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

template <class mint>
mint sum_of_geometric_sequence(mint a, mint r, long long n)
{
    assert(0 < n);

    if (n == 1)
    {
        return a;
    }

    mint x = sum_of_geometric_sequence(a, r, n / 2);

    mint ret = x + r.pow(n / 2) * x;
    if (n & 1)
    {
        ret = (a + r * ret);
    }

    return ret;
};

#line 1 "Math/mod_sum_of_geometric_sequence.hpp"
#include <bits/stdc++.h>
using namespace std;

template <class mint>
mint sum_of_geometric_sequence(mint a, mint r, long long n)
{
    assert(0 < n);

    if (n == 1)
    {
        return a;
    }

    mint x = sum_of_geometric_sequence(a, r, n / 2);

    mint ret = x + r.pow(n / 2) * x;
    if (n & 1)
    {
        ret = (a + r * ret);
    }

    return ret;
};