6.6.6 Polynomial Interpolation

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
114
115
116
117
118
119
120
121

/*

Recovers the unique polynomial of degree less than $n$ that passes through $n$ given points, working
modulo a prime. The $x$ coordinates must be distinct; the $y$ coordinates may be anything. Two
operations are provided: reconstructing the full coefficient vector, and evaluating the
interpolating polynomial at a single point without ever forming the coefficients.

Coefficient recovery uses Newton's divided differences. It first replaces the sample values by their
successive divided differences in place, then expands the Newton basis $1$, $(x - x_0)$,
$(x - x_0)(x - x_1)$, $\ldots$ into ordinary coefficients by multiplying through one factor at a
time. Point evaluation instead applies the Lagrange formula directly, summing each sample weighted
by the product of $(t - x_j) / (x_k - x_j)$ over the other nodes. Coefficient recovery performs
O(n^2) modular inverses, while point evaluation performs $n$.

- `interpolate(x, y)` returns the coefficients $a_0, ..., a_{n-1}$ (constant term first) of the
  polynomial $P$ with $P(`x[i]`) = `y[i]`$, modulo `MOD`. The entries of `x` must be distinct modulo
  `MOD`.
- `interpolate_at(x, y, t)` returns $P(`t`)$ modulo `MOD` for that same polynomial, without building
  its coefficients.

Time Complexity:
- O(n^2 log p) per call to `interpolate()` and O(n^2 + n log p) per call to `interpolate_at()`,
  where $n$ is the number of points and $p$ is `MOD`.

Space Complexity:
- O(n) auxiliary.

*/

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

const int64_t MOD = 998244353;

int64_t powmod(int64_t b, int64_t e, int64_t m) {
  int64_t res = 1;
  for (b %= m; e > 0; e >>= 1) {
    if (e & 1) {
      res = res * b % m;
    }
    b = b * b % m;
  }
  return res;
}

// clang-format off
int64_t modnorm(int64_t a) { return (a %= MOD) < 0 ? a + MOD : a; }
int64_t addmod(int64_t a, int64_t b) { return modnorm(modnorm(a) + modnorm(b)); }
int64_t submod(int64_t a, int64_t b) { return modnorm(modnorm(a) - modnorm(b)); }
int64_t inv(int64_t a) { return powmod(modnorm(a), MOD - 2, MOD); }
// clang-format on

std::vector<int64_t> interpolate(std::vector<int64_t> x, std::vector<int64_t> y) {
  assert(x.size() == y.size());
  int n = static_cast<int>(x.size());
  for (int i = 0; i < n; i++) {
    x[i] = modnorm(x[i]);
    y[i] = modnorm(y[i]);
  }
  // Replace y by its divided differences in place.
  for (int k = 0; k < n; k++) {
    for (int i = k + 1; i < n; i++) {
      y[i] = submod(y[i], y[k]) * inv(submod(x[i], x[k])) % MOD;
    }
  }
  // Expand the Newton basis into ordinary coefficients.
  std::vector<int64_t> res(n), temp(n);
  temp[0] = 1;
  for (int k = 0; k < n; k++) {
    int64_t last = 0;
    for (int i = 0; i < n; i++) {
      res[i] = addmod(res[i], y[k] * temp[i] % MOD);
      int64_t old = temp[i];
      temp[i] = last;
      last = old;
      temp[i] = submod(temp[i], last * x[k] % MOD);
    }
  }
  return res;
}

int64_t interpolate_at(const std::vector<int64_t> &x, const std::vector<int64_t> &y, int64_t t) {
  assert(x.size() == y.size());
  int n = static_cast<int>(x.size());
  t = modnorm(t);
  int64_t res = 0;
  for (int k = 0; k < n; k++) {
    int64_t num = 1, den = 1;
    for (int j = 0; j < n; j++) {
      if (j != k) {
        num = num * submod(t, x[j]) % MOD;
        den = den * submod(x[k], x[j]) % MOD;
      }
    }
    res = addmod(res, modnorm(y[k]) * num % MOD * inv(den) % MOD);
  }
  return res;
}

/*** Example Usage ***/

#include <cassert>
using namespace std;

int main() {
  // P(x) = x^2 + 2x + 3 sampled at x = 0, 1, 2 gives y = 3, 6, 11.
  vector<int64_t> x{0, 1, 2}, y{3, 6, 11};
  assert((interpolate(x, y) == vector<int64_t>{3, 2, 1}));
  assert(interpolate_at(x, y, 5) == 38);  // 25 + 10 + 3.

  // Non-consecutive nodes for P(x) = 2x^3 - x + 4.
  auto poly = [](int64_t t) { return modnorm(2 * t * t * t - t + 4); };
  vector<int64_t> xs{2, 5, 7, 10}, ys;
  for (int64_t v : xs) {
    ys.push_back(poly(v));
  }
  assert((interpolate(xs, ys) == vector<int64_t>{4, MOD - 1, 0, 2}));  // 4 - x + 2x^3.
  assert(interpolate_at(xs, ys, 13) == poly(13));
  return 0;
}