polynomial

TODO : add ntt?

template <typename T>
using Polynomial = std::vector<T>;

namespace polynomial {
  using F = long double;
  using complexF = std::complex<F>;

  constexpr F pi = 3.1415926535897932L;

  void fft(std::vector<complexF>& a) {
    int n = int(a.size());
    for (int i = 1, rev_i = 0; i < n; ++i) {
      for (int bit = n >> 1; ; bit >>= 1) {
        if (((rev_i ^= bit) & bit) != 0) {
          break;
        }
      }
      if (i < rev_i) {
        std::swap(a[i], a[rev_i]);
      }
    }
    for (int b = 1; b < n; b <<= 1) {
      F angle = pi / b;
      complexF w(std::cos(angle), std::sin(angle));
      for (int s = 0; s < n; s += b << 1) {
        complexF theta(1, 0); // for higher precision, use theta(std::cos(angle * i), std::sin(angle * i)) (slower)
        for (int i = 0; i < b; ++i) {
          std::tie(a[s + i], a[s + i + b]) = std::pair {a[s + i] + theta * a[s + i + b], a[s + i] - theta * a[s + i + b]};
          theta *= w;
        }
      }
    }
  }

  void ifft(std::vector<complexF>& a) {
    int n = int(a.size());
    std::reverse(std::next(a.begin()), a.end());
    fft(a);
    for (auto& x : a) {
      x /= n;
    }
  }

  template <typename T>
  Polynomial<T> multiply(const Polynomial<T>& a, const Polynomial<T>& b) {
    if (a.empty() || b.empty()) {
      return {};
    }
    int n = 1;
    while (n < int(a.size()) + int(b.size()) - 1) {
      n <<= 1;
    }
    std::vector<complexF> ca(n);
    std::copy(a.begin(), a.end(), ca.begin());
    std::vector<complexF> cb(n);
    std::copy(b.begin(), b.end(), cb.begin());
    fft(ca); fft(cb);
    for (int i = 0; i < n; ++i) {
      ca[i] *= cb[i];
    }
    ifft(ca);
    Polynomial<T> res(a.size() + b.size() - 1);
    for (int i = 0; i < int(a.size()) + int(b.size()) - 1; ++i) {
      res[i] = std::llround(ca[i].real());
    }
    return res;
  }
} // namespace polynomial

template <typename T>
Polynomial<T> operator+(const Polynomial<T>& a, const Polynomial<T>& b) {
  Polynomial<T> res(std::max(a.size(), b.size()));
  for (int i = 0; i < int(a.size()); ++i) {
    res[i] += a[i];
  }
  for (int i = 0; i < int(b.size()); ++i) {
    res[i] += b[i];
  }
  return res;
}

template <typename T>
Polynomial<T> operator-(const Polynomial<T>& a, const Polynomial<T>& b) {
  Polynomial<T> res(std::max(a.size(), b.size()));
  for (int i = 0; i < int(a.size()); ++i) {
    res[i] += a[i];
  }
  for (int i = 0; i < int(b.size()); ++i) {
    res[i] -= b[i];
  }
  return res;
}

template <typename T>
Polynomial<T>& operator+=(Polynomial<T>& a, const Polynomial<T>& b) {
  a.resize(std::max(a.size(), b.size()));
  for (int i = 0; i < int(b.size()); ++i) {
    a[i] += b[i];
  }
  return a;
}

template <typename T>
Polynomial<T>& operator-=(Polynomial<T>& a, const Polynomial<T>& b) {
  a.resize(std::max(a.size(), b.size()));
  for (int i = 0; i < int(b.size()); ++i) {
    a[i] -= b[i];
  }
  return a;
}

template <typename T>
Polynomial<T> operator*(const Polynomial<T>& a, const Polynomial<T>& b) {
  return polynomial::multiply(a, b);
}

template <typename T>
Polynomial<T>& operator*=(Polynomial<T>& a, const Polynomial<T>& b) {
  return a = a * b;
}