Sum of $k$-th powers of first $n$ positive integers （自然数の $k$ 乗和）
(utilities/kth_power_sum.hpp)

• View this file on GitHub
• Last update: 2024-12-02 06:13:59+09:00
• Include: #include "utilities/kth_power_sum.hpp"

「 $1$ 以上 $n$ 以下の自然数の $k$ 乗の総和」を Faulhaber の公式を用いて計算する．

内部で Bernoulli 数の先頭 $k + 1$ 項を計算する必要がある．このコードでは prefix_sum() メンバ関数が呼ばれた時点で動的に不足している部分を計算し， $\Theta(k^2)$ で動く．より高速にする必要があれば，母関数 $x / (\exp(x) - 1) + x$ （第 1 項の定義に $+1/2$ を採用している点に注意）の冪級数展開を FFT 等で計算してメンバ変数 bernoulli に事前に設定しておけばよい．

使用方法

using mint = ModInt<998244353>;
kth_power_sum<mint> kp;

int k = 1;
long long n = 10;

auto sum = kp.prefix_sum(k, n);  // 1^k + ... + n^k

assert(sum == mint(55));

問題例

• No.665 Bernoulli Bernoulli - yukicoder
• Educational Codeforces Round 7 F. The Sum of the k-th Powers - Codeforces

Verified with

• utilities/test/kth_power_sum.yuki665.test.cpp

Code

#pragma once
#include <vector>

// Computes the sum of the k-th powers
// Complexity:
// - O(k) per query,
// - O(k^2) precomputation (can be reduced to O(k log k) with FFT)
template <class MODINT> struct kth_power_sum {

    // generating function:  x / (e^x - 1) + x
    // NOTE: use B(1) = 1/2 definition
    std::vector<MODINT> bernoulli;

    kth_power_sum() = default;

    void expand() {
        if (bernoulli.empty()) {
            bernoulli = {1};
        } else {
            const int k = bernoulli.size();
            MODINT x(0);
            for (int i = 0; i < k; ++i) { x = -x + bernoulli[i] * MODINT::binom(k + 1, i); }
            bernoulli.push_back(x / (k + 1));
        }
    }

    // Calculate 1^k + 2^k + ... + n^k
    // Based on Faulhaber's formula:
    // S_k(n) = 1 / (k + 1) * sum_{j=0}^{k} B_j * C(k + 1, j) * n^(k + 1 - j)

    template <class T> MODINT prefix_sum(int k, const T &n) {
        while ((int)bernoulli.size() <= k) expand();

        MODINT ret(0), np(1);
        for (int j = k; j >= 0; --j) {
            np *= n;
            ret += MODINT::binom(k + 1, j) * bernoulli[j] * np;
        }
        return ret / (k + 1);
    }
};

#line 2 "utilities/kth_power_sum.hpp"
#include <vector>

// Computes the sum of the k-th powers
// Complexity:
// - O(k) per query,
// - O(k^2) precomputation (can be reduced to O(k log k) with FFT)
template <class MODINT> struct kth_power_sum {

    // generating function:  x / (e^x - 1) + x
    // NOTE: use B(1) = 1/2 definition
    std::vector<MODINT> bernoulli;

    kth_power_sum() = default;

    void expand() {
        if (bernoulli.empty()) {
            bernoulli = {1};
        } else {
            const int k = bernoulli.size();
            MODINT x(0);
            for (int i = 0; i < k; ++i) { x = -x + bernoulli[i] * MODINT::binom(k + 1, i); }
            bernoulli.push_back(x / (k + 1));
        }
    }

    // Calculate 1^k + 2^k + ... + n^k
    // Based on Faulhaber's formula:
    // S_k(n) = 1 / (k + 1) * sum_{j=0}^{k} B_j * C(k + 1, j) * n^(k + 1 - j)

    template <class T> MODINT prefix_sum(int k, const T &n) {
        while ((int)bernoulli.size() <= k) expand();

        MODINT ret(0), np(1);
        for (int j = k; j >= 0; --j) {
            np *= n;
            ret += MODINT::binom(k + 1, j) * bernoulli[j] * np;
        }
        return ret / (k + 1);
    }
};

Back to top page