/
geometric_sequence_sum.hpp
51 lines (46 loc) · 1.21 KB
/
geometric_sequence_sum.hpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
// sum_{i=0}^{n-1}r^i
template <typename T>
T geometic_sequence_sum(T r, ll n) {
if (n == 0) return T(0);
// (r^n, 1+...+r^{n-1})
auto dfs = [&](auto& dfs, ll n) -> pair<T, T> {
if (n == 1) return {r, 1};
auto [x, y] = dfs(dfs, n / 2);
tie(x, y) = mp(x * x, x * y + y);
if (n % 2 == 0) return {x, y};
return {x * r, y + x};
};
return dfs(dfs, n).se;
}
// sum_{i=0}^{n-1}i^kr^i を k=0,1,...,K
template <typename T, int K>
array<T, K + 1> geometic_sequence_sum_K(T r, ll n) {
array<array<T, K + 1>, K + 1> comb{};
comb[0][0] = 1;
FOR(i, K) {
FOR(j, i + 1) {
comb[i + 1][j] += comb[i][j], comb[i + 1][j + 1] += comb[i][j];
}
}
// (n, r^n, sum i^kr^i)
using X = tuple<T, T, array<T, K + 1>>;
auto mul = [&](X& l, X& r) -> X {
auto& [n1, r1, A] = l;
auto& [n2, r2, B] = r;
array<T, K + 1> C;
FOR(k, K + 1) {
C[k] = A[k];
T c = r1;
FOR(j, k + 1) { C[k] += comb[k][j] * c * B[k - j], c *= n1; }
}
return {n1 + n2, r1 * r2, C};
};
X res = {T(0), T(1), array<T, K + 1>{}};
X a = {T(1), T(r), array<T, K + 1>{}};
get<2>(a)[0] = 1;
while (n) {
if (n & 1) res = mul(res, a);
a = mul(a, a), n /= 2;
}
return get<2>(res);
}