exponential_polynomial_sum.hpp¶
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.
Verified by sum_of_exponential_times_polynomial, sum_of_exponential_times_polynomial_limit.
\[
\displaystyle S_d(r,n)=\sum_{i=0}^{n-1} i^d r^i
\]
Implementation¶
#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 degree) {
int count = degree + 1;
std::vector<Mint> values;
values.reserve(count + 1);
if (degree == 0) {
values.assign(count, Mint(1));
return values;
}
std::vector<int> least_prime(count);
std::vector<int> primes;
for (int value = 2; value < count; value++) {
if (least_prime[value] == 0) {
least_prime[value] = value;
primes.push_back(value);
}
for (int prime : primes) {
if (prime > least_prime[value] || prime * std::int64_t(value) >= count) {
break;
}
least_prime[prime * value] = prime;
}
}
values.resize(count);
if (count > 1) {
values[1] = Mint(1);
}
for (int value = 2; value < count; value++) {
if (least_prime[value] == value) {
values[value] = Mint(value).pow(degree);
} else {
int factor = least_prime[value];
values[value] = values[factor] * values[value / factor];
}
}
return values;
}
template <class Mint>
Mint limit_from_values(Mint ratio, int degree,
const std::vector<Mint> &values) {
assert(ratio != Mint(1));
std::vector<Mint> inverse_factorial(degree + 2, Mint(1));
Mint factorial = 1;
for (int value = 1; value <= degree + 1; value++) {
factorial *= Mint(value);
}
inverse_factorial[degree + 1] = Mint(1) / factorial;
for (int value = degree + 1; value > 0; value--) {
inverse_factorial[value - 1] =
inverse_factorial[value] * Mint(value);
}
Mint ratio_power = 1;
Mint reverse_ratio_power = ratio.pow(degree);
Mint inverse_ratio = Mint(1) / ratio;
Mint prefix_sum = 0;
Mint answer = 0;
for (int point = 0; point <= degree; point++) {
prefix_sum += ratio_power * values[point];
Mint term = inverse_factorial[degree - point] *
inverse_factorial[point + 1] * reverse_ratio_power *
prefix_sum;
if ((degree - point) & 1) {
term = -term;
}
answer += term;
ratio_power *= ratio;
reverse_ratio_power *= inverse_ratio;
}
return answer * factorial / (Mint(1) - ratio).pow(degree + 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 ratio, int degree) {
assert(degree >= 0);
assert(degree + 1 < Mint::mod());
assert(ratio != Mint(1));
if (ratio == Mint{}) {
return degree == 0 ? Mint(1) : Mint{};
}
std::vector<Mint> values =
exponential_polynomial_sum_detail::monomial_values<Mint>(degree);
return exponential_polynomial_sum_detail::limit_from_values(
ratio, degree, values);
}
/// @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 ratio, int degree, std::uint64_t count) {
assert(degree >= 0);
assert(degree + 1 < Mint::mod());
if (count == 0) {
return Mint{};
}
if (ratio == Mint{}) {
return degree == 0 ? Mint(1) : Mint{};
}
std::vector<Mint> values =
exponential_polynomial_sum_detail::monomial_values<Mint>(degree);
if (ratio == Mint(1)) {
values.insert(values.begin(), Mint{});
for (int point = 1; point <= degree + 1; point++) {
values[point] += values[point - 1];
}
return lagrange_consecutive(values, count);
}
Mint limit = exponential_polynomial_sum_detail::limit_from_values(
ratio, degree, values);
Mint prefix_sum = 0;
Mint ratio_power = 1;
Mint inverse_ratio_power = 1;
Mint inverse_ratio = Mint(1) / ratio;
values.push_back(Mint{});
for (int point = 0; point <= degree; point++) {
Mint monomial = values[point];
values[point] = inverse_ratio_power * (prefix_sum - limit);
prefix_sum += ratio_power * monomial;
ratio_power *= ratio;
inverse_ratio_power *= inverse_ratio;
}
values[degree + 1] = inverse_ratio_power * (prefix_sum - limit);
return limit + ratio.pow(count) * lagrange_consecutive(values, count);
}
} // 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 < values.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> &values, Integer x) {
static_assert(std::is_integral_v<Integer>);
assert(!values.empty());
if constexpr (std::is_signed_v<Integer>) {
if (x >= 0 && std::uint64_t(x) < values.size()) {
return values[std::size_t(x)];
}
} else if (x < values.size()) {
return values[std::size_t(x)];
}
int n = int(values.size());
T point = T(x);
std::vector<T> prefix(n + 1, T(1));
std::vector<T> suffix(n + 1, T(1));
for (int index = 0; index < n; index++) {
prefix[index + 1] = prefix[index] * (point - T(index));
}
for (int index = n - 1; index >= 0; index--) {
suffix[index] = suffix[index + 1] * (point - T(index));
}
std::vector<T> inverse_factorial(n, T(1));
T factorial = T(1);
for (int value = 1; value < n; value++) {
factorial *= T(value);
}
inverse_factorial[n - 1] = T(1) / factorial;
for (int value = n - 1; value >= 1; value--) {
inverse_factorial[value - 1] = inverse_factorial[value] * T(value);
}
T result{};
for (int index = 0; index < n; index++) {
T coefficient = prefix[index] * suffix[index + 1] *
inverse_factorial[index] * inverse_factorial[n - 1 - index];
if ((n - 1 - index) & 1) {
coefficient = -coefficient;
}
result += values[index] * coefficient;
}
return result;
}
/// @brief Return sum_{i=1}^n i^exponent over a field in O(exponent).
template <class T> T power_sum(std::uint64_t n, int exponent) {
assert(exponent >= 0);
auto power = [&](T value, int degree) {
T result = T(1);
while (degree > 0) {
if (degree & 1) {
result *= value;
}
value *= value;
degree >>= 1;
}
return result;
};
std::vector<T> values(exponent + 2);
for (int point = 1; point < int(values.size()); point++) {
values[point] = values[point - 1] + power(T(point), exponent);
}
return lagrange_consecutive(values, n);
}
} // namespace noya
namespace noya {
namespace exponential_polynomial_sum_detail {
template <class Mint>
std::vector<Mint> monomial_values(int degree) {
int count = degree + 1;
std::vector<Mint> values;
values.reserve(count + 1);
if (degree == 0) {
values.assign(count, Mint(1));
return values;
}
std::vector<int> least_prime(count);
std::vector<int> primes;
for (int value = 2; value < count; value++) {
if (least_prime[value] == 0) {
least_prime[value] = value;
primes.push_back(value);
}
for (int prime : primes) {
if (prime > least_prime[value] || prime * std::int64_t(value) >= count) {
break;
}
least_prime[prime * value] = prime;
}
}
values.resize(count);
if (count > 1) {
values[1] = Mint(1);
}
for (int value = 2; value < count; value++) {
if (least_prime[value] == value) {
values[value] = Mint(value).pow(degree);
} else {
int factor = least_prime[value];
values[value] = values[factor] * values[value / factor];
}
}
return values;
}
template <class Mint>
Mint limit_from_values(Mint ratio, int degree,
const std::vector<Mint> &values) {
assert(ratio != Mint(1));
std::vector<Mint> inverse_factorial(degree + 2, Mint(1));
Mint factorial = 1;
for (int value = 1; value <= degree + 1; value++) {
factorial *= Mint(value);
}
inverse_factorial[degree + 1] = Mint(1) / factorial;
for (int value = degree + 1; value > 0; value--) {
inverse_factorial[value - 1] =
inverse_factorial[value] * Mint(value);
}
Mint ratio_power = 1;
Mint reverse_ratio_power = ratio.pow(degree);
Mint inverse_ratio = Mint(1) / ratio;
Mint prefix_sum = 0;
Mint answer = 0;
for (int point = 0; point <= degree; point++) {
prefix_sum += ratio_power * values[point];
Mint term = inverse_factorial[degree - point] *
inverse_factorial[point + 1] * reverse_ratio_power *
prefix_sum;
if ((degree - point) & 1) {
term = -term;
}
answer += term;
ratio_power *= ratio;
reverse_ratio_power *= inverse_ratio;
}
return answer * factorial / (Mint(1) - ratio).pow(degree + 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 ratio, int degree) {
assert(degree >= 0);
assert(degree + 1 < Mint::mod());
assert(ratio != Mint(1));
if (ratio == Mint{}) {
return degree == 0 ? Mint(1) : Mint{};
}
std::vector<Mint> values =
exponential_polynomial_sum_detail::monomial_values<Mint>(degree);
return exponential_polynomial_sum_detail::limit_from_values(
ratio, degree, values);
}
/// @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 ratio, int degree, std::uint64_t count) {
assert(degree >= 0);
assert(degree + 1 < Mint::mod());
if (count == 0) {
return Mint{};
}
if (ratio == Mint{}) {
return degree == 0 ? Mint(1) : Mint{};
}
std::vector<Mint> values =
exponential_polynomial_sum_detail::monomial_values<Mint>(degree);
if (ratio == Mint(1)) {
values.insert(values.begin(), Mint{});
for (int point = 1; point <= degree + 1; point++) {
values[point] += values[point - 1];
}
return lagrange_consecutive(values, count);
}
Mint limit = exponential_polynomial_sum_detail::limit_from_values(
ratio, degree, values);
Mint prefix_sum = 0;
Mint ratio_power = 1;
Mint inverse_ratio_power = 1;
Mint inverse_ratio = Mint(1) / ratio;
values.push_back(Mint{});
for (int point = 0; point <= degree; point++) {
Mint monomial = values[point];
values[point] = inverse_ratio_power * (prefix_sum - limit);
prefix_sum += ratio_power * monomial;
ratio_power *= ratio;
inverse_ratio_power *= inverse_ratio;
}
values[degree + 1] = inverse_ratio_power * (prefix_sum - limit);
return limit + ratio.pow(count) * lagrange_consecutive(values, count);
}
} // namespace noya