Made lagrange a class.

Author: a1fredbao

src/alfred/math/lagrange.hpp

@@ -2,26 +2,78 @@
 #define AFMT_LAGRANGE
 
 #include "comb.hpp"
+#include "vec-inv.hpp"
 #include <cassert>
 #include <iostream>
 #include <vector>
 
 // TODO: write lagrange as a class, supporting: O(n^2) init, O(n) query
 
 template <class mint>
-inline mint lagrange(std::vector<mint> x, std::vector<mint> y, mint k) {
-    mint ans = 0, cur;
-    const int n = x.size();
-    for (int i = 0; i < n; i++) {
-        cur = y[i];
-        for (int j = 0; j < n; j++) {
-            if (j == i) continue;
-            cur *= (k - x[j]) / (x[i] - x[j]);
+class Lagrange {
+private:
+    std::vector<mint> x, y, b;
+
+public:
+    Lagrange(void) = default;
+    Lagrange(std::vector<mint> _x, std::vector<mint> _y) : x(_x), y(_y) {
+        for (size_t i = 0; i < x.size(); i++) insert(x[i], y[i]);
+    }
+    inline void insert(mint x0, mint y0) {
+        b.push_back(y0);
+        std::vector<mint> tmp(x.size());
+        for (size_t i = 0; i < x.size(); i++) {
+            tmp[i] = x0 - x[i];
+        }
+        VecInv<mint> inv(tmp);
+        for (size_t i = 0; i < x.size(); i++) {
+            b.back() *= inv[i], b[i] *= -inv[i];
         }
-        ans += cur;
+        x.push_back(x0), y.push_back(y0);
     }
-    return ans;
-}
+    inline mint query(mint k) {
+        mint ans = 0, tot = 1;
+        const int n = x.size();
+        std::vector<mint> tmp(n);
+        for (int i = 0; i < n; i++) {
+            if (x[i] == k) return y[i];
+            tmp[i] = k - x[i], tot *= tmp[i];
+        }
+        VecInv<mint> inv(tmp);
+        for (int i = 0; i < n; i++) {
+            ans += b[i] * tot * inv[i];
+        }
+        return ans;
+    }
+    std::vector<mint> coefficient(void) { // now O(n^2), TODO: maintain it dynamically.
+        int n = x.size(), i;
+        // F(k) = \prod (k - x_i): degree = n, n + 1 coefficients.
+        std::vector<mint> F(n + 1);
+        for (i = 0, F[0] = 1; i < n; i++) {
+            for (int j = i + 1; j >= 0; j--) {
+                F[j] *= -x[i];
+                if (j) F[j] += F[j - 1];
+            }
+        }
+        mint delta, c;
+        std::vector<mint> ans(n), res(n);
+        auto div = [&](mint xi) {
+            delta = 0;
+            for (int i = n; i > 0; i--) {
+                res[i - 1] = F[i] + delta;
+                delta = (F[i] + delta) * xi;
+            }
+        };
+        for (int i = 0; i < n; i++) {
+            c = b[i], div(x[i]);
+            for (int j = 0; j < n; j++) {
+                ans[j] += c * res[j];
+            }
+        }
+        std::reverse(ans.begin(), ans.end());
+        return ans;
+    }
+};
 
 // y[0] is placeholder.
 // If for all integer x_i in [1, n], we have f(x_i) = y_i (mod p), find f(k) mod p.
@@ -52,41 +104,4 @@ inline mint sum_of_kth_powers(mint n, int k) {
     return cont_lagrange(Y, n);
 }
 
-template <class mint>
-std::vector<mint> find_coefficient(
-    std::vector<mint> x, std::vector<mint> y
-) {
-    // F(k) = \prod (k - x_i): n degree, n + 1 coefficients.
-    int n = x.size(), i;
-    std::vector<mint> F(n + 1);
-    assert(n == (int)y.size());
-    for (i = 0, F[0] = 1; i < n; i++) {
-        for (int j = i + 1; j >= 0; j--) {
-            F[j] *= -x[i];
-            if (j) F[j] += F[j - 1];
-        }
-    }
-    mint delta, c;
-    std::vector<mint> ans(n), res(n);
-    auto div = [&](mint xi) {
-        delta = 0;
-        for (int i = n; i > 0; i--) {
-            res[i - 1] = F[i] + delta;
-            delta = (F[i] + delta) * xi;
-        }
-        return res;
-    };
-    for (int i = 0; i < n; i++) {
-        c = y[i];
-        for (int j = 0; j < n; j++) {
-            if (i != j) c /= x[i] - x[j];
-        }
-        div(x[i]);
-        for (int j = 0; j < n; j++) {
-            ans[j] += c * res[j];
-        }
-    }
-    return ans;
-}
-
 #endif // AFMT_LAGRANGE