6.6.5 Linear Recurrence (Kitamasa)

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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

/*

Computes the $n$-th term of a linear recurrence modulo a prime in logarithmic time. Given a
recurrence of order $L$, $s_i = c_0 s_{i-1} + c_1 s_{i-2} + \ldots + c_{L-1} s_{i-L}$, together with
the first $L$ terms, the naive approach unrolls $n$ steps. Kitamasa's method instead jumps directly
to index $n$ in O(L^2 log n) by working with the characteristic polynomial. The coefficient layout
matches the output of Berlekamp-Massey, so the two compose directly: guess the recurrence from
sampled values, then jump to any index.

The key identity is that the $n$-th term is a fixed linear combination of the first $L$ terms, with
weights given by $x^n$ reduced modulo the characteristic polynomial
$f(x) = x^L - c_0 x^{L-1} - ... - c_{L-1}$. The reduction $x^L = c_0 x^{L-1} + ... + c_{L-1}$ lets
any product of two degree-$<L$ polynomials be folded back to degree $<L$, so $x^n \mod f$ is
obtained by binary exponentiation, and the $n$-th term is the dot product of its coefficients with
the initial terms.

- `kth_term(rec, init, n)` returns $s_n$ modulo `MOD`, where `rec` holds the coefficients
  $c_0, ..., c_{L-1}$ and `init` holds at least $L$ initial terms $s_0, ..., s_{L-1}$. Indexing is
  0-based and `n` $\geq 0$. Values are reduced modulo `MOD` internally.

Time Complexity:
- O(L^2 log n) per call.

Space Complexity:
- O(L) auxiliary.

*/

#include <cassert>
#include <cstdint>
#include <vector>

const int64_t MOD = 998244353;

std::vector<int64_t> combine(
    const std::vector<int64_t> &a, const std::vector<int64_t> &b, const std::vector<int64_t> &c
) {
  std::vector<int64_t> res(a.size() + b.size() - 1);
  for (int i = 0; i < static_cast<int>(a.size()); i++) {
    for (int j = 0; j < static_cast<int>(b.size()); j++) {
      res[i + j] = (res[i + j] + a[i] * b[j]) % MOD;
    }
  }
  int L = static_cast<int>(c.size());
  for (int d = static_cast<int>(res.size()) - 1; d >= L; d--) {
    if (res[d] != 0) {
      for (int j = 0; j < L; j++) {
        res[d - 1 - j] = (res[d - 1 - j] + res[d] * c[j]) % MOD;
      }
    }
  }
  res.resize(L);
  return res;
}

int64_t kth_term(const std::vector<int64_t> &rec, const std::vector<int64_t> &init, int64_t n) {
  assert(n >= 0 && init.size() >= rec.size());
  std::vector<int64_t> c(rec);
  for (int64_t &x : c) {
    x = (x % MOD + MOD) % MOD;
  }
  int L = static_cast<int>(c.size());
  if (n < static_cast<int64_t>(init.size())) {
    return ((init[n] % MOD) + MOD) % MOD;
  }
  if (L == 0) {
    return 0;
  }
  std::vector<int64_t> result(L), base(L);
  result[0] = 1;  // The polynomial 1 = x^0.
  if (L == 1) {
    base[0] = c[0];  // x is congruent to c_0 modulo f.
  } else {
    base[1] = 1;  // The polynomial x.
  }
  for (int64_t e = n; e > 0; e >>= 1) {
    if (e & 1) {
      result = combine(result, base, c);
    }
    base = combine(base, base, c);
  }
  int64_t ans = 0;
  for (int j = 0; j < L; j++) {
    ans = (ans + result[j] * (((init[j] % MOD) + MOD) % MOD)) % MOD;
  }
  return ans;
}

/*** Example Usage ***/

using namespace std;

int main() {
  // Fibonacci s_i = s_{i-1} + s_{i-2} with s_0 = s_1 = 1: 1, 1, 2, 3, 5, 8, ...
  vector<int64_t> fib_rec{1, 1}, fib_init{1, 1};
  assert(kth_term(fib_rec, fib_init, 0) == 1);
  assert(kth_term(fib_rec, fib_init, 9) == 55);
  assert(kth_term(fib_rec, fib_init, 10) == 89);

  // Tribonacci s_i = s_{i-1} + s_{i-2} + s_{i-3} with s_0,s_1,s_2 = 0,1,1.
  vector<int64_t> trib_rec{1, 1, 1}, trib_init{0, 1, 1};
  assert(kth_term(trib_rec, trib_init, 8) == 44);

  // Geometric s_i = 3 s_{i-1}, s_0 = 1: matches 3^n modulo MOD.
  vector<int64_t> geo_rec{3}, geo_init{1};
  int64_t pw = 1;
  for (int i = 0; i < 50; i++) {
    assert(kth_term(geo_rec, geo_init, i) == pw);
    pw = pw * 3 % MOD;
  }
  assert(kth_term({-1}, {1}, 5) == MOD - 1);
  return 0;
}