Skip to content

sparse_formal_power_series.hpp

SECTIONMath INCLUDEnoya/sparse_formal_power_series.hpp

Invert a sparse formal power series. From f g = 1, each new coefficient of g is a dot product against the nonconstant terms of f.

Verified by exp_of_formal_power_series_sparse, inv_of_formal_power_series_sparse, log_of_formal_power_series_sparse, pow_of_formal_power_series_sparse, sqrt_of_formal_power_series_sparse.

\[ \displaystyle f(x)g(x)=1 \]

Implementation

View on GitHub

#ifndef NOYA_SPARSE_FORMAL_POWER_SERIES_HPP
#define NOYA_SPARSE_FORMAL_POWER_SERIES_HPP 1

/// @complexity Time: O(nk), where k is the number of nonzero input terms.
/// Space: O(n + k).

#include "noya/mod_sqrt.hpp"

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <optional>
#include <utility>
#include <vector>

namespace noya {

template <class Mint>
using sparse_power_series = std::vector<std::pair<int, Mint>>;

namespace sparse_fps_detail {

template <class Mint>
std::vector<Mint> coefficient_inverses(int size) {
  std::vector<Mint> inverse(size);
  if (size > 1) {
    inverse[1] = Mint(1);
  }
  for (int value = 2; value < size; value++) {
    inverse[value] = -Mint(Mint::mod() / value) * inverse[Mint::mod() % value];
  }
  return inverse;
}

template <class Mint>
void check_terms(int size, const sparse_power_series<Mint> &terms) {
  int previous = -1;
  for (auto [degree, coefficient] : terms) {
    assert(previous < degree && degree < size && coefficient != Mint{});
    previous = degree;
  }
}

template <class Mint>
std::vector<Mint> unit_power(int size,
                             const sparse_power_series<Mint> &terms,
                             Mint exponent) {
  std::vector<Mint> result(size);
  if (size == 0) return result;
  result[0] = Mint(1);
  std::vector<Mint> inverse = coefficient_inverses<Mint>(size);
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree == 0) continue;
      if (term_degree > degree) break;
      value += coefficient * result[degree - term_degree] *
               (exponent * Mint(term_degree) - Mint(degree - term_degree));
    }
    result[degree] = value * inverse[degree];
  }
  return result;
}

template <class Mint>
Mint scalar_power(Mint value, std::uint64_t exponent) {
  Mint result = Mint(1);
  while (exponent > 0) {
    if (exponent & 1) result *= value;
    value *= value;
    exponent >>= 1;
  }
  return result;
}

} // namespace sparse_fps_detail

/// @brief Invert a sparse formal power series.  From `f g = 1`, each new
/// coefficient of g is a dot product against the nonconstant terms of f.
template <class Mint>
std::vector<Mint> sparse_fps_inverse(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  assert(!terms.empty() && terms.front().first == 0);
  std::vector<Mint> result(size);
  if (size == 0) return result;
  Mint inverse_constant = Mint(1) / terms.front().second;
  result[0] = inverse_constant;
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree == 0) continue;
      if (term_degree > degree) break;
      value += coefficient * result[degree - term_degree];
    }
    result[degree] = -value * inverse_constant;
  }
  return result;
}

/// @brief Compute log(f) for sparse f with f(0)=1.  Comparing coefficients in
/// `f (log f)' = f'` gives one linear recurrence per output coefficient.
template <class Mint>
std::vector<Mint> sparse_fps_logarithm(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  assert(!terms.empty() && terms.front().first == 0 &&
         terms.front().second == Mint(1));
  std::vector<Mint> result(size);
  std::vector<Mint> inverse =
      sparse_fps_detail::coefficient_inverses<Mint>(size);
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree == 0) continue;
      if (term_degree > degree) break;
      if (term_degree == degree) {
        value += Mint(degree) * coefficient;
      } else {
        value -= coefficient * Mint(degree - term_degree) *
                 result[degree - term_degree];
      }
    }
    result[degree] = value * inverse[degree];
  }
  return result;
}

/// @brief Compute exp(f) for sparse f with f(0)=0.  The differential identity
/// `g' = f' g` gives the recurrence directly.
template <class Mint>
std::vector<Mint> sparse_fps_exponential(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  assert(terms.empty() || terms.front().first > 0);
  std::vector<Mint> result(size);
  if (size == 0) return result;
  result[0] = Mint(1);
  std::vector<Mint> inverse =
      sparse_fps_detail::coefficient_inverses<Mint>(size);
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree > degree) break;
      value += Mint(term_degree) * coefficient *
               result[degree - term_degree];
    }
    result[degree] = value * inverse[degree];
  }
  return result;
}

/// @brief Raise a sparse series to a nonnegative integer exponent.  The first
/// nonzero monomial determines the output shift and scale; after normalization
/// `f g' = exponent f' g` yields an O(nk) recurrence.
template <class Mint>
std::vector<Mint> sparse_fps_power(
    int size, const sparse_power_series<Mint> &terms,
    std::uint64_t exponent) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  std::vector<Mint> result(size);
  if (size == 0) return result;
  if (exponent == 0) {
    result[0] = Mint(1);
    return result;
  }
  if (terms.empty()) return result;
  int first_degree = terms.front().first;
  if (first_degree > 0 &&
      exponent > std::uint64_t((size - 1) / first_degree)) {
    return result;
  }
  int shift = int(std::uint64_t(first_degree) * exponent);
  int target = size - shift;
  Mint leading = terms.front().second;
  sparse_power_series<Mint> normalized;
  normalized.reserve(terms.size());
  for (auto [degree, coefficient] : terms) {
    int normalized_degree = degree - first_degree;
    if (normalized_degree >= target) break;
    normalized.emplace_back(normalized_degree, coefficient / leading);
  }
  std::vector<Mint> unit = sparse_fps_detail::unit_power(
      target, normalized, Mint(exponent));
  Mint scale = sparse_fps_detail::scalar_power(leading, exponent);
  for (int index = 0; index < target; index++) {
    result[shift + index] = scale * unit[index];
  }
  return result;
}

/// @brief Return one square root of a sparse series when it exists.  Its first
/// degree must be even and its leading coefficient a quadratic residue; after
/// removing both, the unit series is raised to the field exponent 1/2.
template <class Mint>
std::optional<std::vector<Mint>> sparse_fps_square_root(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  if (terms.empty()) return std::vector<Mint>(size);
  int first_degree = terms.front().first;
  if (first_degree % 2 != 0) return std::nullopt;
  auto root = mod_sqrt(std::uint64_t(terms.front().second.val()),
                       std::uint64_t(Mint::mod()));
  if (!root) return std::nullopt;
  int shift = first_degree / 2;
  int target = size - first_degree;
  Mint leading = terms.front().second;
  sparse_power_series<Mint> normalized;
  normalized.reserve(terms.size());
  for (auto [degree, coefficient] : terms) {
    int normalized_degree = degree - first_degree;
    if (normalized_degree >= target) break;
    normalized.emplace_back(normalized_degree, coefficient / leading);
  }
  std::vector<Mint> unit = sparse_fps_detail::unit_power(
      target, normalized, Mint(1) / Mint(2));
  std::vector<Mint> result(size);
  for (int index = 0; index < target; index++) {
    result[shift + index] = Mint(*root) * unit[index];
  }
  return result;
}

} // namespace noya

#endif // NOYA_SPARSE_FORMAL_POWER_SERIES_HPP
#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <numeric>
#include <optional>
#include <utility>
#include <vector>

/// @complexity Time: O(nk), where k is the number of nonzero input terms.
/// Space: O(n + k).

/// @complexity Time: O(log^2 p).
/// Space: O(1).

/// @complexity Time: O(log^3 n) primality testing; Pollard-rho factorization is expected about O(n^(1/4)).
/// Space: O(log n) recursion and factors.

namespace noya {
namespace factorize_internal {

using u64 = std::uint64_t;
using u128 = unsigned __int128;

inline u64 multiply_mod(u64 a, u64 b, u64 mod) {
  return u64(u128(a) * b % mod);
}

inline u64 power_mod(u64 a, u64 exponent, u64 mod) {
  u64 result = 1;
  while (exponent > 0) {
    if (exponent & 1) {
      result = multiply_mod(result, a, mod);
    }
    a = multiply_mod(a, a, mod);
    exponent >>= 1;
  }
  return result;
}

inline bool miller_rabin(u64 n) {
  if (n < 2) {
    return false;
  }
  for (u64 p :
       std::array<u64, 12>{2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37}) {
    if (n % p == 0) {
      return n == p;
    }
  }
  int shift = __builtin_ctzll(n - 1);
  u64 odd = (n - 1) >> shift;
  for (u64 base :
       std::array<u64, 7>{2, 325, 9375, 28178, 450775, 9780504, 1795265022}) {
    if (base % n == 0) {
      continue;
    }
    u64 value = power_mod(base % n, odd, n);
    if (value == 1 || value == n - 1) {
      continue;
    }
    bool composite = true;
    for (int i = 1; i < shift; i++) {
      value = multiply_mod(value, value, n);
      if (value == n - 1) {
        composite = false;
        break;
      }
    }
    if (composite) {
      return false;
    }
  }
  return true;
}

inline u64 splitmix64(u64 &state) {
  u64 z = (state += 0x9e3779b97f4a7c15ULL);
  z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
  z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
  return z ^ (z >> 31);
}

inline u64 pollard_rho(u64 n) {
  if (n % 2 == 0) {
    return 2;
  }
  if (n % 3 == 0) {
    return 3;
  }
  static u64 state = 0x123456789abcdef0ULL;
  while (true) {
    u64 y = splitmix64(state) % (n - 1) + 1;
    u64 c = splitmix64(state) % (n - 1) + 1;
    constexpr u64 block = 128;
    u64 g = 1;
    u64 r = 1;
    u64 q = 1;
    u64 x = 0;
    u64 saved_y = 0;
    auto next = [&](u64 value) {
      return u64((u128(multiply_mod(value, value, n)) + c) % n);
    };
    while (g == 1) {
      x = y;
      for (u64 i = 0; i < r; i++) {
        y = next(y);
      }
      for (u64 offset = 0; offset < r && g == 1; offset += block) {
        saved_y = y;
        for (u64 i = 0; i < std::min(block, r - offset); i++) {
          y = next(y);
          u64 difference = x > y ? x - y : y - x;
          q = multiply_mod(q, difference, n);
        }
        g = std::gcd(q, n);
      }
      r <<= 1;
    }
    if (g == n) {
      do {
        saved_y = next(saved_y);
        u64 difference = x > saved_y ? x - saved_y : saved_y - x;
        g = std::gcd(difference, n);
      } while (g == 1);
    }
    if (g != n) {
      return g;
    }
  }
}

inline void collect_factors(u64 n, std::vector<u64> &result) {
  if (n == 1) {
    return;
  }
  if (miller_rabin(n)) {
    result.push_back(n);
    return;
  }
  u64 factor = pollard_rho(n);
  collect_factors(factor, result);
  collect_factors(n / factor, result);
}

} // namespace factorize_internal

/// @brief Deterministic Miller-Rabin primality test for unsigned 64-bit
/// integers.
inline bool is_prime(std::uint64_t n) {
  return factorize_internal::miller_rabin(n);
}

/// @brief Return the prime factors of n with multiplicity in increasing order.
inline std::vector<std::uint64_t> prime_factors(std::uint64_t n) {
  assert(n >= 1);
  std::vector<std::uint64_t> result;
  factorize_internal::collect_factors(n, result);
  std::sort(result.begin(), result.end());
  return result;
}

/// @brief Return the prime factorization of n as (prime, exponent) pairs.
inline std::vector<std::pair<std::uint64_t, int>> factorize(std::uint64_t n) {
  std::vector<std::pair<std::uint64_t, int>> result;
  for (std::uint64_t p : prime_factors(n)) {
    if (result.empty() || result.back().first != p) {
      result.emplace_back(p, 1);
    } else {
      result.back().second++;
    }
  }
  return result;
}

} // namespace noya

namespace noya {

/// @brief Compute the smaller square root modulo a prime, or nullopt if no
/// square root exists.
inline std::optional<std::uint64_t> mod_sqrt(std::uint64_t value,
                                             std::uint64_t modulus) {
  assert(modulus >= 2 && is_prime(modulus));
  value %= modulus;
  if (modulus == 2 || value == 0) {
    return value;
  }
  using factorize_internal::multiply_mod;
  using factorize_internal::power_mod;
  if (power_mod(value, (modulus - 1) / 2, modulus) != 1) {
    return std::nullopt;
  }
  if (modulus % 4 == 3) {
    std::uint64_t root = power_mod(value, (modulus + 1) / 4, modulus);
    return std::min(root, modulus - root);
  }

  std::uint64_t odd = modulus - 1;
  int exponent = 0;
  while ((odd & 1) == 0) {
    odd >>= 1;
    exponent++;
  }
  std::uint64_t non_residue = 2;
  while (power_mod(non_residue, (modulus - 1) / 2, modulus) != modulus - 1) {
    non_residue++;
  }

  std::uint64_t root = power_mod(value, (odd + 1) / 2, modulus);
  std::uint64_t remainder = power_mod(value, odd, modulus);
  std::uint64_t step = power_mod(non_residue, odd, modulus);
  int remaining = exponent;
  while (remainder != 1) {
    std::uint64_t squared = remainder;
    int shift = 0;
    while (squared != 1 && shift < remaining) {
      squared = multiply_mod(squared, squared, modulus);
      shift++;
    }
    assert(shift < remaining);
    std::uint64_t multiplier =
        power_mod(step, std::uint64_t(1) << (remaining - shift - 1), modulus);
    root = multiply_mod(root, multiplier, modulus);
    step = multiply_mod(multiplier, multiplier, modulus);
    remainder = multiply_mod(remainder, step, modulus);
    remaining = shift;
  }
  return std::min(root, modulus - root);
}

} // namespace noya

namespace noya {

template <class Mint>
using sparse_power_series = std::vector<std::pair<int, Mint>>;

namespace sparse_fps_detail {

template <class Mint>
std::vector<Mint> coefficient_inverses(int size) {
  std::vector<Mint> inverse(size);
  if (size > 1) {
    inverse[1] = Mint(1);
  }
  for (int value = 2; value < size; value++) {
    inverse[value] = -Mint(Mint::mod() / value) * inverse[Mint::mod() % value];
  }
  return inverse;
}

template <class Mint>
void check_terms(int size, const sparse_power_series<Mint> &terms) {
  int previous = -1;
  for (auto [degree, coefficient] : terms) {
    assert(previous < degree && degree < size && coefficient != Mint{});
    previous = degree;
  }
}

template <class Mint>
std::vector<Mint> unit_power(int size,
                             const sparse_power_series<Mint> &terms,
                             Mint exponent) {
  std::vector<Mint> result(size);
  if (size == 0) return result;
  result[0] = Mint(1);
  std::vector<Mint> inverse = coefficient_inverses<Mint>(size);
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree == 0) continue;
      if (term_degree > degree) break;
      value += coefficient * result[degree - term_degree] *
               (exponent * Mint(term_degree) - Mint(degree - term_degree));
    }
    result[degree] = value * inverse[degree];
  }
  return result;
}

template <class Mint>
Mint scalar_power(Mint value, std::uint64_t exponent) {
  Mint result = Mint(1);
  while (exponent > 0) {
    if (exponent & 1) result *= value;
    value *= value;
    exponent >>= 1;
  }
  return result;
}

} // namespace sparse_fps_detail

/// @brief Invert a sparse formal power series.  From `f g = 1`, each new
/// coefficient of g is a dot product against the nonconstant terms of f.
template <class Mint>
std::vector<Mint> sparse_fps_inverse(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  assert(!terms.empty() && terms.front().first == 0);
  std::vector<Mint> result(size);
  if (size == 0) return result;
  Mint inverse_constant = Mint(1) / terms.front().second;
  result[0] = inverse_constant;
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree == 0) continue;
      if (term_degree > degree) break;
      value += coefficient * result[degree - term_degree];
    }
    result[degree] = -value * inverse_constant;
  }
  return result;
}

/// @brief Compute log(f) for sparse f with f(0)=1.  Comparing coefficients in
/// `f (log f)' = f'` gives one linear recurrence per output coefficient.
template <class Mint>
std::vector<Mint> sparse_fps_logarithm(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  assert(!terms.empty() && terms.front().first == 0 &&
         terms.front().second == Mint(1));
  std::vector<Mint> result(size);
  std::vector<Mint> inverse =
      sparse_fps_detail::coefficient_inverses<Mint>(size);
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree == 0) continue;
      if (term_degree > degree) break;
      if (term_degree == degree) {
        value += Mint(degree) * coefficient;
      } else {
        value -= coefficient * Mint(degree - term_degree) *
                 result[degree - term_degree];
      }
    }
    result[degree] = value * inverse[degree];
  }
  return result;
}

/// @brief Compute exp(f) for sparse f with f(0)=0.  The differential identity
/// `g' = f' g` gives the recurrence directly.
template <class Mint>
std::vector<Mint> sparse_fps_exponential(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  assert(terms.empty() || terms.front().first > 0);
  std::vector<Mint> result(size);
  if (size == 0) return result;
  result[0] = Mint(1);
  std::vector<Mint> inverse =
      sparse_fps_detail::coefficient_inverses<Mint>(size);
  for (int degree = 1; degree < size; degree++) {
    Mint value{};
    for (auto [term_degree, coefficient] : terms) {
      if (term_degree > degree) break;
      value += Mint(term_degree) * coefficient *
               result[degree - term_degree];
    }
    result[degree] = value * inverse[degree];
  }
  return result;
}

/// @brief Raise a sparse series to a nonnegative integer exponent.  The first
/// nonzero monomial determines the output shift and scale; after normalization
/// `f g' = exponent f' g` yields an O(nk) recurrence.
template <class Mint>
std::vector<Mint> sparse_fps_power(
    int size, const sparse_power_series<Mint> &terms,
    std::uint64_t exponent) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  std::vector<Mint> result(size);
  if (size == 0) return result;
  if (exponent == 0) {
    result[0] = Mint(1);
    return result;
  }
  if (terms.empty()) return result;
  int first_degree = terms.front().first;
  if (first_degree > 0 &&
      exponent > std::uint64_t((size - 1) / first_degree)) {
    return result;
  }
  int shift = int(std::uint64_t(first_degree) * exponent);
  int target = size - shift;
  Mint leading = terms.front().second;
  sparse_power_series<Mint> normalized;
  normalized.reserve(terms.size());
  for (auto [degree, coefficient] : terms) {
    int normalized_degree = degree - first_degree;
    if (normalized_degree >= target) break;
    normalized.emplace_back(normalized_degree, coefficient / leading);
  }
  std::vector<Mint> unit = sparse_fps_detail::unit_power(
      target, normalized, Mint(exponent));
  Mint scale = sparse_fps_detail::scalar_power(leading, exponent);
  for (int index = 0; index < target; index++) {
    result[shift + index] = scale * unit[index];
  }
  return result;
}

/// @brief Return one square root of a sparse series when it exists.  Its first
/// degree must be even and its leading coefficient a quadratic residue; after
/// removing both, the unit series is raised to the field exponent 1/2.
template <class Mint>
std::optional<std::vector<Mint>> sparse_fps_square_root(
    int size, const sparse_power_series<Mint> &terms) {
  assert(size >= 0);
  sparse_fps_detail::check_terms(size, terms);
  if (terms.empty()) return std::vector<Mint>(size);
  int first_degree = terms.front().first;
  if (first_degree % 2 != 0) return std::nullopt;
  auto root = mod_sqrt(std::uint64_t(terms.front().second.val()),
                       std::uint64_t(Mint::mod()));
  if (!root) return std::nullopt;
  int shift = first_degree / 2;
  int target = size - first_degree;
  Mint leading = terms.front().second;
  sparse_power_series<Mint> normalized;
  normalized.reserve(terms.size());
  for (auto [degree, coefficient] : terms) {
    int normalized_degree = degree - first_degree;
    if (normalized_degree >= target) break;
    normalized.emplace_back(normalized_degree, coefficient / leading);
  }
  std::vector<Mint> unit = sparse_fps_detail::unit_power(
      target, normalized, Mint(1) / Mint(2));
  std::vector<Mint> result(size);
  for (int index = 0; index < target; index++) {
    result[shift + index] = Mint(*root) * unit[index];
  }
  return result;
}

} // namespace noya