Skip to content

exponential_polynomial_sum.hpp

SECTIONMath INCLUDEnoya/exponential_polynomial_sum.hpp

计算多项式乘指数项的有限和或前缀和,例如求和 \(P(i)r^i\)

\[ \displaystyle S_d(r,n)=\sum_{i=0}^{n-1} i^d r^i \]

Complexity: Time: O(d log log(d + 2) + log n). Space: O(d).

AC 记录:sum_of_exponential_times_polynomial, sum_of_exponential_times_polynomial_limit

跳到代码 · GitHub ↗

Implementation

当前头文件,省略 include guard;依赖见 #include

/// @complexity Time: O(d log log(d + 2) + log n). Space: O(d).

#include "noya/lagrange_interpolation.hpp"

#include <cassert>
#include <cstdint>
#include <vector>

namespace noya {

namespace exponential_polynomial_sum_detail {

template <class Mint> std::vector<Mint> monomial_values(int deg) {
  int cnt = deg + 1;
  std::vector<Mint> vs;
  vs.reserve(cnt + 1);
  if (deg == 0) {
    vs.assign(cnt, Mint(1));
    return vs;
  }

  std::vector<int> mnp(cnt);
  std::vector<int> ps;
  for (int val = 2; val < cnt; val++) {
    if (mnp[val] == 0) {
      mnp[val] = val;
      ps.push_back(val);
    }
    for (int p : ps) {
      if (p > mnp[val] || p * std::int64_t(val) >= cnt) {
        break;
      }
      mnp[p * val] = p;
    }
  }

  vs.resize(cnt);
  if (cnt > 1) {
    vs[1] = Mint(1);
  }
  for (int val = 2; val < cnt; val++) {
    if (mnp[val] == val) {
      vs[val] = Mint(val).pow(deg);
    } else {
      int fct = mnp[val];
      vs[val] = vs[fct] * vs[val / fct];
    }
  }
  return vs;
}

template <class Mint>
Mint limit_from_values(Mint r, int deg, const std::vector<Mint> &vs) {
  assert(r != Mint(1));
  std::vector<Mint> ifc(deg + 2, Mint(1));
  Mint fac = 1;
  for (int val = 1; val <= deg + 1; val++) {
    fac *= Mint(val);
  }
  ifc[deg + 1] = Mint(1) / fac;
  for (int val = deg + 1; val > 0; val--) {
    ifc[val - 1] = ifc[val] * Mint(val);
  }

  Mint rp = 1;
  Mint rrp = r.pow(deg);
  Mint ir = Mint(1) / r;
  Mint pre = 0;
  Mint ans = 0;
  for (int poi = 0; poi <= deg; poi++) {
    pre += rp * vs[poi];
    Mint ter = ifc[deg - poi] * ifc[poi + 1] * rrp * pre;
    if ((deg - poi) & 1) {
      ter = -ter;
    }
    ans += ter;
    rp *= r;
    rrp *= ir;
  }
  return ans * fac / (Mint(1) - r).pow(deg + 1);
}

} // namespace exponential_polynomial_sum_detail

/// @brief Return sum_{i>=0} r^i i^d over a finite field for r != 1, with
/// 0^0=1. A degree-d polynomial is determined by its values at 0..d. Writing
/// it in the consecutive-point Lagrange basis and summing each basis polynomial
/// against the geometric series reduces the answer to factorial weights and
/// prefix sums of r^i i^d.
template <class Mint> Mint exponential_polynomial_sum_limit(Mint r, int deg) {
  assert(deg >= 0);
  assert(deg + 1 < Mint::mod());
  assert(r != Mint(1));
  if (r == Mint{}) {
    return deg == 0 ? Mint(1) : Mint{};
  }
  std::vector<Mint> vs =
      exponential_polynomial_sum_detail::monomial_values<Mint>(deg);
  return exponential_polynomial_sum_detail::limit_from_values(r, deg, vs);
}

/// @brief Return sum_{i=0}^{n-1} r^i i^d, with 0^0=1. For r=1 the prefix
/// sum is a degree-(d+1) polynomial in n. Otherwise let C be the infinite
/// weighted sum; r^{-n}(S(n)-C) is a degree-d polynomial, so d+2 sampled
/// prefix sums and one consecutive-point interpolation determine S(n), even
/// for very large n.
template <class Mint>
Mint exponential_polynomial_sum(Mint r, int deg, std::uint64_t cnt) {
  assert(deg >= 0);
  assert(deg + 1 < Mint::mod());
  if (cnt == 0) {
    return Mint{};
  }
  if (r == Mint{}) {
    return deg == 0 ? Mint(1) : Mint{};
  }

  std::vector<Mint> vs =
      exponential_polynomial_sum_detail::monomial_values<Mint>(deg);
  if (r == Mint(1)) {
    vs.insert(vs.begin(), Mint{});
    for (int poi = 1; poi <= deg + 1; poi++) {
      vs[poi] += vs[poi - 1];
    }
    return lagrange_consecutive(vs, cnt);
  }

  Mint lim = exponential_polynomial_sum_detail::limit_from_values(r, deg, vs);
  Mint pre = 0;
  Mint rp = 1;
  Mint irp = 1;
  Mint ir = Mint(1) / r;
  vs.push_back(Mint{});
  for (int poi = 0; poi <= deg; poi++) {
    Mint mon = vs[poi];
    vs[poi] = irp * (pre - lim);
    pre += rp * mon;
    rp *= r;
    irp *= ir;
  }
  vs[deg + 1] = irp * (pre - lim);
  return lim + r.pow(cnt) * lagrange_consecutive(vs, cnt);
}

} // namespace noya
#ifndef NOYA_EXPONENTIAL_POLYNOMIAL_SUM_HPP
#define NOYA_EXPONENTIAL_POLYNOMIAL_SUM_HPP 1

/// @complexity Time: O(d log log(d + 2) + log n). Space: O(d).

#include "noya/lagrange_interpolation.hpp"

#include <cassert>
#include <cstdint>
#include <vector>

namespace noya {

namespace exponential_polynomial_sum_detail {

template <class Mint> std::vector<Mint> monomial_values(int deg) {
  int cnt = deg + 1;
  std::vector<Mint> vs;
  vs.reserve(cnt + 1);
  if (deg == 0) {
    vs.assign(cnt, Mint(1));
    return vs;
  }

  std::vector<int> mnp(cnt);
  std::vector<int> ps;
  for (int val = 2; val < cnt; val++) {
    if (mnp[val] == 0) {
      mnp[val] = val;
      ps.push_back(val);
    }
    for (int p : ps) {
      if (p > mnp[val] || p * std::int64_t(val) >= cnt) {
        break;
      }
      mnp[p * val] = p;
    }
  }

  vs.resize(cnt);
  if (cnt > 1) {
    vs[1] = Mint(1);
  }
  for (int val = 2; val < cnt; val++) {
    if (mnp[val] == val) {
      vs[val] = Mint(val).pow(deg);
    } else {
      int fct = mnp[val];
      vs[val] = vs[fct] * vs[val / fct];
    }
  }
  return vs;
}

template <class Mint>
Mint limit_from_values(Mint r, int deg, const std::vector<Mint> &vs) {
  assert(r != Mint(1));
  std::vector<Mint> ifc(deg + 2, Mint(1));
  Mint fac = 1;
  for (int val = 1; val <= deg + 1; val++) {
    fac *= Mint(val);
  }
  ifc[deg + 1] = Mint(1) / fac;
  for (int val = deg + 1; val > 0; val--) {
    ifc[val - 1] = ifc[val] * Mint(val);
  }

  Mint rp = 1;
  Mint rrp = r.pow(deg);
  Mint ir = Mint(1) / r;
  Mint pre = 0;
  Mint ans = 0;
  for (int poi = 0; poi <= deg; poi++) {
    pre += rp * vs[poi];
    Mint ter = ifc[deg - poi] * ifc[poi + 1] * rrp * pre;
    if ((deg - poi) & 1) {
      ter = -ter;
    }
    ans += ter;
    rp *= r;
    rrp *= ir;
  }
  return ans * fac / (Mint(1) - r).pow(deg + 1);
}

} // namespace exponential_polynomial_sum_detail

/// @brief Return sum_{i>=0} r^i i^d over a finite field for r != 1, with
/// 0^0=1. A degree-d polynomial is determined by its values at 0..d. Writing
/// it in the consecutive-point Lagrange basis and summing each basis polynomial
/// against the geometric series reduces the answer to factorial weights and
/// prefix sums of r^i i^d.
template <class Mint> Mint exponential_polynomial_sum_limit(Mint r, int deg) {
  assert(deg >= 0);
  assert(deg + 1 < Mint::mod());
  assert(r != Mint(1));
  if (r == Mint{}) {
    return deg == 0 ? Mint(1) : Mint{};
  }
  std::vector<Mint> vs =
      exponential_polynomial_sum_detail::monomial_values<Mint>(deg);
  return exponential_polynomial_sum_detail::limit_from_values(r, deg, vs);
}

/// @brief Return sum_{i=0}^{n-1} r^i i^d, with 0^0=1. For r=1 the prefix
/// sum is a degree-(d+1) polynomial in n. Otherwise let C be the infinite
/// weighted sum; r^{-n}(S(n)-C) is a degree-d polynomial, so d+2 sampled
/// prefix sums and one consecutive-point interpolation determine S(n), even
/// for very large n.
template <class Mint>
Mint exponential_polynomial_sum(Mint r, int deg, std::uint64_t cnt) {
  assert(deg >= 0);
  assert(deg + 1 < Mint::mod());
  if (cnt == 0) {
    return Mint{};
  }
  if (r == Mint{}) {
    return deg == 0 ? Mint(1) : Mint{};
  }

  std::vector<Mint> vs =
      exponential_polynomial_sum_detail::monomial_values<Mint>(deg);
  if (r == Mint(1)) {
    vs.insert(vs.begin(), Mint{});
    for (int poi = 1; poi <= deg + 1; poi++) {
      vs[poi] += vs[poi - 1];
    }
    return lagrange_consecutive(vs, cnt);
  }

  Mint lim = exponential_polynomial_sum_detail::limit_from_values(r, deg, vs);
  Mint pre = 0;
  Mint rp = 1;
  Mint irp = 1;
  Mint ir = Mint(1) / r;
  vs.push_back(Mint{});
  for (int poi = 0; poi <= deg; poi++) {
    Mint mon = vs[poi];
    vs[poi] = irp * (pre - lim);
    pre += rp * mon;
    rp *= r;
    irp *= ir;
  }
  vs[deg + 1] = irp * (pre - lim);
  return lim + r.pow(cnt) * lagrange_consecutive(vs, cnt);
}

} // namespace noya

#endif // NOYA_EXPONENTIAL_POLYNOMIAL_SUM_HPP
#include <cassert>
#include <cstdint>
#include <type_traits>
#include <vector>

/// @complexity Time: O(d log log(d + 2) + log n). Space: O(d).

/// @complexity Time: O(n) per evaluation at consecutive points.
/// Space: O(n).

namespace noya {

/// @brief Evaluate the degree < vs.size() polynomial known at consecutive
/// points 0,1,... in O(n) over a field.
template <class T, class Integer>
T lagrange_consecutive(const std::vector<T> &vs, Integer x) {
  static_assert(std::is_integral_v<Integer>);
  assert(!vs.empty());
  if constexpr (std::is_signed_v<Integer>) {
    if (x >= 0 && std::uint64_t(x) < vs.size()) {
      return vs[std::size_t(x)];
    }
  } else if (x < vs.size()) {
    return vs[std::size_t(x)];
  }
  int n = int(vs.size());
  T poi = T(x);
  std::vector<T> pre(n + 1, T(1));
  std::vector<T> suf(n + 1, T(1));
  for (int idx = 0; idx < n; idx++) {
    pre[idx + 1] = pre[idx] * (poi - T(idx));
  }
  for (int idx = n - 1; idx >= 0; idx--) {
    suf[idx] = suf[idx + 1] * (poi - T(idx));
  }
  std::vector<T> ifc(n, T(1));
  T fac = T(1);
  for (int val = 1; val < n; val++) {
    fac *= T(val);
  }
  ifc[n - 1] = T(1) / fac;
  for (int val = n - 1; val >= 1; val--) {
    ifc[val - 1] = ifc[val] * T(val);
  }
  T res{};
  for (int idx = 0; idx < n; idx++) {
    T cf = pre[idx] * suf[idx + 1] * ifc[idx] * ifc[n - 1 - idx];
    if ((n - 1 - idx) & 1) {
      cf = -cf;
    }
    res += vs[idx] * cf;
  }
  return res;
}

/// @brief Return sum_{i=1}^n i^exp over a field in O(exp).
template <class T> T power_sum(std::uint64_t n, int exp) {
  assert(exp >= 0);
  auto pw = [&](T val, int deg) {
    T res = T(1);
    while (deg > 0) {
      if (deg & 1) {
        res *= val;
      }
      val *= val;
      deg >>= 1;
    }
    return res;
  };
  std::vector<T> vs(exp + 2);
  for (int poi = 1; poi < int(vs.size()); poi++) {
    vs[poi] = vs[poi - 1] + pw(T(poi), exp);
  }
  return lagrange_consecutive(vs, n);
}

} // namespace noya

namespace noya {

namespace exponential_polynomial_sum_detail {

template <class Mint> std::vector<Mint> monomial_values(int deg) {
  int cnt = deg + 1;
  std::vector<Mint> vs;
  vs.reserve(cnt + 1);
  if (deg == 0) {
    vs.assign(cnt, Mint(1));
    return vs;
  }

  std::vector<int> mnp(cnt);
  std::vector<int> ps;
  for (int val = 2; val < cnt; val++) {
    if (mnp[val] == 0) {
      mnp[val] = val;
      ps.push_back(val);
    }
    for (int p : ps) {
      if (p > mnp[val] || p * std::int64_t(val) >= cnt) {
        break;
      }
      mnp[p * val] = p;
    }
  }

  vs.resize(cnt);
  if (cnt > 1) {
    vs[1] = Mint(1);
  }
  for (int val = 2; val < cnt; val++) {
    if (mnp[val] == val) {
      vs[val] = Mint(val).pow(deg);
    } else {
      int fct = mnp[val];
      vs[val] = vs[fct] * vs[val / fct];
    }
  }
  return vs;
}

template <class Mint>
Mint limit_from_values(Mint r, int deg, const std::vector<Mint> &vs) {
  assert(r != Mint(1));
  std::vector<Mint> ifc(deg + 2, Mint(1));
  Mint fac = 1;
  for (int val = 1; val <= deg + 1; val++) {
    fac *= Mint(val);
  }
  ifc[deg + 1] = Mint(1) / fac;
  for (int val = deg + 1; val > 0; val--) {
    ifc[val - 1] = ifc[val] * Mint(val);
  }

  Mint rp = 1;
  Mint rrp = r.pow(deg);
  Mint ir = Mint(1) / r;
  Mint pre = 0;
  Mint ans = 0;
  for (int poi = 0; poi <= deg; poi++) {
    pre += rp * vs[poi];
    Mint ter = ifc[deg - poi] * ifc[poi + 1] * rrp * pre;
    if ((deg - poi) & 1) {
      ter = -ter;
    }
    ans += ter;
    rp *= r;
    rrp *= ir;
  }
  return ans * fac / (Mint(1) - r).pow(deg + 1);
}

} // namespace exponential_polynomial_sum_detail

/// @brief Return sum_{i>=0} r^i i^d over a finite field for r != 1, with
/// 0^0=1. A degree-d polynomial is determined by its values at 0..d. Writing
/// it in the consecutive-point Lagrange basis and summing each basis polynomial
/// against the geometric series reduces the answer to factorial weights and
/// prefix sums of r^i i^d.
template <class Mint> Mint exponential_polynomial_sum_limit(Mint r, int deg) {
  assert(deg >= 0);
  assert(deg + 1 < Mint::mod());
  assert(r != Mint(1));
  if (r == Mint{}) {
    return deg == 0 ? Mint(1) : Mint{};
  }
  std::vector<Mint> vs =
      exponential_polynomial_sum_detail::monomial_values<Mint>(deg);
  return exponential_polynomial_sum_detail::limit_from_values(r, deg, vs);
}

/// @brief Return sum_{i=0}^{n-1} r^i i^d, with 0^0=1. For r=1 the prefix
/// sum is a degree-(d+1) polynomial in n. Otherwise let C be the infinite
/// weighted sum; r^{-n}(S(n)-C) is a degree-d polynomial, so d+2 sampled
/// prefix sums and one consecutive-point interpolation determine S(n), even
/// for very large n.
template <class Mint>
Mint exponential_polynomial_sum(Mint r, int deg, std::uint64_t cnt) {
  assert(deg >= 0);
  assert(deg + 1 < Mint::mod());
  if (cnt == 0) {
    return Mint{};
  }
  if (r == Mint{}) {
    return deg == 0 ? Mint(1) : Mint{};
  }

  std::vector<Mint> vs =
      exponential_polynomial_sum_detail::monomial_values<Mint>(deg);
  if (r == Mint(1)) {
    vs.insert(vs.begin(), Mint{});
    for (int poi = 1; poi <= deg + 1; poi++) {
      vs[poi] += vs[poi - 1];
    }
    return lagrange_consecutive(vs, cnt);
  }

  Mint lim = exponential_polynomial_sum_detail::limit_from_values(r, deg, vs);
  Mint pre = 0;
  Mint rp = 1;
  Mint irp = 1;
  Mint ir = Mint(1) / r;
  vs.push_back(Mint{});
  for (int poi = 0; poi <= deg; poi++) {
    Mint mon = vs[poi];
    vs[poi] = irp * (pre - lim);
    pre += rp * mon;
    rp *= r;
    irp *= ir;
  }
  vs[deg + 1] = irp * (pre - lim);
  return lim + r.pow(cnt) * lagrange_consecutive(vs, cnt);
}

} // namespace noya