More inplace stuff

Author: adamant-pwn

cp-algo/math/poly.hpp

@@ -33,9 +33,25 @@ namespace cp_algo::math {
         poly_t operator + (poly_t const& t) const {return poly_t(*this) += t;}
         poly_t operator - (poly_t const& t) const {return poly_t(*this) -= t;}
 
-        poly_t mod_xk(size_t k) const {return poly::impl::mod_xk(*this, k);} // %= x^k
-        poly_t mul_xk(size_t k) const {return poly::impl::mul_xk(*this, k);} // *= x^k
-        poly_t div_xk(size_t k) const {return poly::impl::div_xk(*this, k);} // /= x^k
+        poly_t& mod_xk_inplace(size_t k) {
+            a.resize(std::min(size(a), k));
+            normalize();
+            return *this;
+        }
+        poly_t& mul_xk_inplace(size_t k) {
+            a.insert(begin(a), k, T(0));
+            normalize();
+            return *this;
+        }
+        poly_t& div_xk_inplace(size_t k) {
+            a.erase(begin(a), min(k, size(a)));
+            normalize();
+            return *this;
+        }
+        poly_t mod_xk(size_t k) const {return poly_t(*this).mod_xk_inplace(k);}
+        poly_t mul_xk(size_t k) const {return poly_t(*this).mul_xk_inplace(k);}
+        poly_t div_xk(size_t k) const {return poly_t(*this).div_xk_inplace(k);}
+        
         poly_t substr(size_t l, size_t k) const {return poly::impl::substr(*this, l, k);}
 
         poly_t operator *= (const poly_t &t) {fft::mul(a, t.a); normalize(); return *this;}
@@ -135,17 +151,28 @@ namespace cp_algo::math {
         bool operator == (const poly_t &t) const {return a == t.a;}
         bool operator != (const poly_t &t) const {return a != t.a;}
 
-        poly_t deriv(int k = 1) const { // calculate derivative
+        poly_t& deriv_inplace(int k = 1) {
             if(deg() + 1 < k) {
-                return poly_t(T(0));
+                return *this = poly_t{};
             }
-            std::vector<T> res(deg() + 1 - k);
             for(int i = k; i <= deg(); i++) {
-                res[i - k] = fact<T>(i) * rfact<T>(i - k) * a[i];
+                a[i - k] = fact<T>(i) * rfact<T>(i - k) * a[i];
             }
-            return res;
+            a.resize(deg() + 1 - k);
+            return *this;
+        }
+        poly_t deriv(int k = 1) const { // calculate derivative
+            return poly_t(*this).deriv_inplace(k);
+        }
+
+        poly_t& integr_inplace() {
+            a.push_back(0);
+            for(int i = deg() - 1; i >= 0; i--) {
+                a[i + 1] = a[i] * small_inv<T>(i + 1);
+            }
+            a[0] = 0;
+            return *this;
         }
-        
         poly_t integr() const { // calculate integral with C = 0
             std::vector<T> res(deg() + 2);
             for(int i = 0; i <= deg(); i++) {
@@ -165,9 +192,15 @@ namespace cp_algo::math {
             return res;
         }
 
-        poly_t log(size_t n) const { // calculate log p(x) mod x^n
+        // calculate log p(x) mod x^n
+        poly_t log_inplace(size_t n) {
             assert(a[0] == T(1));
-            return (mod_xk(n).deriv() * inv(n)).mod_xk(n - 1).integr();
+            auto t = mod_xk_inplace(n).inv(n);
+            deriv_inplace();
+            return (*this *= t).mod_xk_inplace(n - 1).integr_inplace();
+        }
+        poly_t log(size_t n) const {
+            return poly_t(*this).log_inplace(n);
         }
 
         poly_t exp(size_t n) const { // calculate exp p(x) mod x^n
@@ -585,8 +618,7 @@ namespace cp_algo::math {
 
         // inverse series mod x^n
         poly_t& inv_inplace(size_t n) {
-            poly::impl::inv_inplace(*this, n);
-            return *this;
+            return poly::impl::inv_inplace(*this, n);
         }
         poly_t inv(size_t n) const {
             return poly_t(*this).inv_inplace(n);

cp-algo/math/poly/impl/base.hpp

@@ -34,24 +34,6 @@ namespace cp_algo::math::poly::impl {
         normalize(p);
         return p;
     }
-    auto mod_xk(auto const& p, size_t k) {
-        return std::vector(begin(p.a), begin(p.a) + std::min(k, p.a.size()));
-    }
-    auto mul_xk(auto p, int k) {
-        if(k < 0) {
-            return p.div_xk(-k);
-        }
-        p.a.insert(begin(p.a), k, 0);
-        normalize(p);
-        return p;
-    }
-    template<typename poly>
-    poly div_xk(poly const& p, int k) {
-        if(k < 0) {
-            return p.mul_xk(-k);
-        }
-        return std::vector(begin(p.a) + std::min<size_t>(k, p.a.size()), end(p.a));
-    }
     auto substr(auto const& p, size_t l, size_t k) {
         return std::vector(
             begin(p.a) + std::min(l, p.a.size()),

cp-algo/math/poly/impl/div.hpp

@@ -91,12 +91,11 @@ namespace cp_algo::math::poly::impl {
         ).div_xk(2*q0.deg()).mod_xk(n);
     }
     template<typename poly>
-    void inv_inplace(poly&& p, size_t n) {
+    poly& inv_inplace(poly& p, size_t n) {
         using poly_t = std::decay_t<poly>;
         using base = poly_t::base;
         if(n == 1) {
-            p = base(1) / p[0];
-            return;
+            return p = base(1) / p[0];
         }
         // Q(-x) = P0(x^2) + xP1(x^2)
         auto [q0, q1] = p.bisect(n);
@@ -107,7 +106,8 @@ namespace cp_algo::math::poly::impl {
         auto q1f = fft::dft<base>(q1.a, N);
 
         // Q(x)*Q(-x) = Q0(x^2)^2 - x^2 Q1(x^2)^2
-        auto qq = poly_t(q0f * q0f) - poly_t(q1f * q1f).mul_xk(1);
+        auto qq = poly_t(q0f * q0f) - poly_t(q1f * q1f).mul_xk_inplace(1);
+
         inv_inplace(qq, (n + 1) / 2);
         auto qqf = fft::dft<base>(qq.a, N);
 
@@ -120,6 +120,8 @@ namespace cp_algo::math::poly::impl {
             p.a[i + 1] = -B[i / 2];
         }
         p.a.pop_back();
+        p.normalize();
+        return p;
     }
 }
 #endif // CP_ALGO_MATH_POLY_IMPL_DIV_HPP

verify/poly/log.test.cpp

@@ -16,7 +16,7 @@ void solve() {
     cin >> n;
     vector<base> a(n);
     copy_n(istream_iterator<base>(cin), n, begin(a));
-    polyn(a).log(n).print(n);
+    polyn(a).log_inplace(n).print(n);
 }
 
 signed main() {