Lagrange Interpolation

There are some well-known formulas

$$\sum_{i=1}^n i = 1 + 2 + \dots + n = \frac{n * (n + 1)}{2}$$

$$\sum_{i=1}^n i^2 = 1^2 + 2^2 + \dots + n^2 = \frac{n * (n + 1) * (2n + 1)}{6}$$

$$\sum_{i=1}^n i^3 = 1^3 + 2^3 + \dots + n^3 = (\frac{n * (n + 1)}{2}) ^ 2$$

Then what is the value of the following sum of the k-th power?

$$\sum_{i=1}^n i^k = 1^k + 2^k + \dots + n^k \mod 10^9 + 7$$

given $1 <= n <= 10^9$ and $0 <= k <= 10^6$

The target sum will be a degree $(k + 1)$ polynomial and we can interpolate the answer with $(k + 2)$ data points, i.e. $degree(f) + 1$ points. In order to find the data points, we need to calculate $f(0) = 0$, $f(x) = f(x - 1) + x ^ k$.

If there is less than $k + 2$ data points, we can calculate the answer directly.

if (x <= k + 1) {
  int s = 0;
  for (int i = 1; i <= x; i++) {
    s = (s + qpow(i, k)) % mod;
  }
  return s;
}

Otherwise, let's say $f(x_1) = y_1, f(x_2) = y_2, ..., f(x_n) = y_n$ and $f$ is the unique $(n - 1)$ degree polynomial and we are interested in $f(x) = \sum_{i=1}^n f_i(x)$ where $f_i(x) = y_i * \prod_{j=1, j!=i}^n \frac{x - x_j}{x_i - x_j}$. Therefore, we have our Lagrange interpolation as $f(x) = \sum_{i=1}^n y^i \prod_{j=1, j!=i}^n \frac{x - x_j}{x_i - x_j}$.

However, we need $O(n^2)$ to calculate $f(x)$. Let's substitute $x_i = i$ and $y_i = f(i)$ and we got $f(x) = \sum_{i=1}^n f(i) \frac{\prod_{j=1, j!=i}^n x - j}{\prod_{j=1, j!=i}^n i - j}$. What can we do for numerator and denonminator here? For numerator, we can per-calculate the prefix and suffix product of $x$, for each $i$, we can calculate the nubmerator in $O(1)$.

$$\prod_{j=1, j!=i}^n x - x_j = [(x - 1)(x - 2)...(x-(i-1))] * [(x - (i + 1))*(x - (i + 2))...(x - n)]$$

vector<int> pre(k + 2), suf(k + 2);
pre[0] = x;
suf[k + 1] = x - (k + 1);
for (int i = 1; i <= k; i++) pre[i] = pre[i - 1] * (x - i) % mod;
for (int i = k; i >= 1; i--) suf[i] = suf[i + 1] * (x - i) % mod;

For denominator, we can precompute the factorials using their inverse in $O(1)$ also.

$$\prod_{j=1, j!=i}^n i - j = [(i - 1)(i - 2)(i - 3)...(i - (i - 1))] * [i - (i + 1)(i - (i + 2)...(i - n)] = (-1)^{n - i} (n - i)!(i - 1)!$$

int qpow(int base, int exp) {
  int res = 1;
  while (exp) {
    if (exp & 1) res = (res * base) % mod;
    base = (base * base) % mod;
    exp >>= 1;
  }
  return res;
}

unordered_map<int, int> rv_m;
int rv(int x) {
  if (rv_m.count(x)) {
    return rv_m[x];
  }
  return rv_m[x] = qpow(x, mod - 2);
}

vector<int> inv(k + 2);
inv[0] = 1;
for (int i = 1; i <= k + 1; i++) inv[i] = inv[i - 1] * rv(i) % mod;

Overall, we can calculate $f(x)$ in $O(n)$.

Complete Code:

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

#define int long long
const int mod = 1e9 + 7;

int qpow(int base, int exp) {
  int res = 1;
  while (exp) {
    if (exp & 1) res = (res * base) % mod;
    base = (base * base) % mod;
    exp >>= 1;
  }
  return res;
}

unordered_map<int, int> rv_m;
int rv(int x) {
  if (rv_m.count(x)) {
    return rv_m[x];
  }
  return rv_m[x] = qpow(x, mod - 2);
}

int lagrange_interpolate(int x, int k, bool bf = false) {
  if (k == 0) return x;
  // find 1 ^ k + 2 ^ k + ... + x ^ k
  // (k + 1) degree polynomial -> (k + 2) points
  if (x <= k + 1 || bf) {
    int s = 0;
    for (int i = 1; i <= x; i++) {
      s = (s + qpow(i, k)) % mod;
    }
    return s;
  }
  vector<int> pre(k + 2), suf(k + 2), inv(k + 2);
  inv[0] = 1, pre[0] = x;
  suf[k + 1] = x - (k + 1);
  for (int i = 1; i <= k; i++) pre[i] = pre[i - 1] * (x - i) % mod;
  for (int i = k; i >= 1; i--) suf[i] = suf[i + 1] * (x - i) % mod;
  for (int i = 1; i <= k + 1; i++) inv[i] = inv[i - 1] * rv(i) % mod;
  int ans = 0;
  int yi = 0;  // 0 ^ k + ~ i ^ k
  int num, denom;
  for (int i = 0; i <= k + 1; i++) {
    yi = (yi + qpow(i, k)) % mod; // interpolate point: (i, yi)
    if (i == 0) num = suf[1];
    else if (i == k + 1) num = pre[k];
    else num = pre[i - 1] * suf[i + 1] % mod;  // numerator
    denom = inv[i] * inv[k + 1 - i] % mod;  // denominator
    if ((i + k) & 1) ans += (yi * num % mod) * denom % mod;
    else ans -= (yi * num % mod) * denom % mod;
    ans = (ans % mod + mod) % mod;
  }
  return ans;
}

void solve() {
  int n, k; 
  cin >> n >> k;
  cout << lagrange_interpolate(n, k) << endl;
}

int32_t main() {
  ios_base::sync_with_stdio(false), cin.tie(nullptr);

  // int T; cin >> T;
  // while(T--) solve();
  solve();
  return 0;
}