等比数列の和

math/geometric_sum.hpp · GitHub で見る

共通ページ

公比 \(r\) の等比数列の先頭 \(n\) 項の和を求めます。除算を使わないので、\(r=1\) や、\(r-1\) の逆元が存在しない法にも使えます。

#include "lib/math/geometric_sum.hpp"
#include "lib/math/mint.hpp"
#include <cassert>

assert(geometric_sum(3LL, 4) == 40);       // 1 + 3 + 9 + 27
assert(geometric_sum(2LL, 3LL, 4) == 80);  // 2 + 6 + 18 + 54
assert(geometric_sum(-2LL, 4) == -5);      // 1 - 2 + 4 - 8
assert(geometric_sum(1LL, 100) == 100);
assert(geometric_sum(0LL, 0) == 0);

const mint s = geometric_sum(mint(3), 1000000000000000000ULL);
assert(s * mint(2) == mint(3).pow(1000000000000000000LL) - mint(1));

関数	内容・返り値
geometric_sum(ratio, n)	\(\displaystyle\sum_{i=0}^{n-1}\mathrm{ratio}^i\) を返す。ratio の型 T が返り値の型
geometric_sum(first, ratio, n)	\(\mathrm{first}\cdot\displaystyle\sum_{i=0}^{n-1}\mathrm{ratio}^i\) を返す。first と ratio は同じ型 T

n は bool 以外の組込み整数型、__int128_t、__uint128_t で指定します。0 項の和は T(0)、負の n は std::invalid_argument を送出します。既定引数はありません。

T には T(0)・T(1) による構築、コピー、加算 +、乗算 * が必要です。加算・乗算は結合則と分配則を満たし、0 と 1 はそれぞれ零元と単位元とします。逆元や減算・除算は必要ありません。自前の mint、ACL の static/dynamic modint、合成数を法とする modint、多倍長整数にも使えます。

T が bool 以外の組込み整数型または 128 bit 整数型なら、途中の加算・乗算で型の範囲を超えた場合に std::overflow_error を送出します。他の型の演算の安全性はその型に従います。型を広げたい場合は geometric_sum<__int128_t>(ratio, n) のように指定できます。型の演算が対応していれば constexpr でも使えます。

加算・乗算を定数時間として、時間は \(O(\log n)\)、追加空間は \(O(1)\) です。n = 0 は \(O(1)\) 時間です。

依存ライブラリ

• 整数演算の共通部品

コード

#ifndef GEOMETRIC_SUM_HPP
#define GEOMETRIC_SUM_HPP

#include "../components/integer_arithmetic.hpp"
#include <limits>
#include <stdexcept>

template <class T, class N>
	requires integer_arithmetic_detail::integer<N>
constexpr T geometric_sum(T ratio, N n) {
	using integer_arithmetic_detail::add, integer_arithmetic_detail::multiply;
	if constexpr (std::numeric_limits<N>::is_signed) if (n < 0) throw std::invalid_argument("geometric_sum length");
	T sum = T(0), power = T(1), block = T(1);
	while (n) {
		if (n & 1) sum = add(sum, multiply(power, block));
		if (n == 1) break;
		if (n & 1) power = multiply(power, ratio);
		block = add(block, multiply(ratio, block));
		if (n >= 4) ratio = multiply(ratio, ratio);
		n >>= 1;
	}
	return sum;
}

template <class T, class N>
	requires integer_arithmetic_detail::integer<N>
constexpr T geometric_sum(T first, T ratio, N n) {
	return integer_arithmetic_detail::multiply(first, geometric_sum(ratio, n));
}

#endif

Back to top page