cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:heavy_check_mark: 乗法的関数の prefix sum

math/multiplicative_prefix_sum.hpp · GitHub で見る

共通ページ

乗法的関数 \(f\) の \(\displaystyle F(x)=\sum_{i=1}^x f(i)\) を計算します。f(1) = 1 とし、互いに素な \(a,b\) について \(f(ab)=f(a)f(b)\) を満たす関数が対象です。上限を含み、F(0) = 0 です。

例えば、オイラーのトーティエント関数の和は次のように求められます。

#include "lib/math/multiplicative_prefix_sum.hpp"
#include "lib/math/mint.hpp"

MultiplicativePrefixSum<mint> sieve(100);
auto prime_prefix = sieve.prime_sum_linear(mint(-1), mint(1));
auto phi = [](std::uint64_t p, unsigned, std::uint64_t pe) {
    return mint(pe - pe / p);
};

auto answer = sieve.sum(prime_prefix, phi); // 3044
auto prefix = sieve.prefix_sum(prime_prefix, phi);
auto at10 = sieve.get(prefix, 10); // 32

利用には、素数上の prefix sum の表 prime_prefix と、素数冪に対する値を返す関数 prime_power が必要です。

prime_prefix に渡すもの

prime_prefix は長さ sieve.size() の std::vector<T> です。q = sieve.quotients() とすると、各要素に次の値を入れます。

\[\displaystyle \texttt{prime\_prefix}[i]=G(q_i)=\sum_{\substack{p\leq q_i\\p\text{ prime}}} f(p).\]

quotients()[i] 以下の素数だけについて、f(p) を合計した値です。1 や合成数の寄与は含めません。添字 i は quotients() の添字で、素数の番号や整数値そのものではありません。全体の累積和への復元は prefix_sum が行い、f(1) = 1 と高い素数冪・異なる素数の積の寄与を加えます。

例えば n = 10 で、冒頭と同じトーティエント関数を使う場合は、素数上の値が \(f(p)=p-1\) なので、渡す表は次のようになります。

i quotients()[i] 合計する値 prime_prefix[i]
0 10 (2-1)+(3-1)+(5-1)+(7-1) 13
1 5 (2-1)+(3-1)+(5-1) 7
2 3 (2-1)+(3-1) 3
3 2 2-1 1
4 1 素数なし 0

つまり、この場合の引数は std::vector<T>{13, 7, 3, 1, 0} です。prime_power(p,e,pe) に T(pe-pe/p) を返す関数を渡すと、prefix_sum の返す表は {32, 10, 4, 2, 1} になります。

素数上の値が \(f(p)=a+bp\) なら、sieve.prime_sum_linear(a,b) の戻り値をそのまま渡せます。最初のトーティエントの例では \(f(p)=p-1\) なので、prime_sum_linear(T(-1),T(1)) で作ります。他の形では prime_sum で作った表の線形結合や、自分で用意した正しい表を渡せます。同じ上限 n の quotients() に合わせた表を使ってください。

API

以下では \(q_i=\texttt{quotients()}[i]\) とします。

操作 内容・返り値
MultiplicativePrefixSum<T>(n = 0) 上限 n の計算領域を作る
max_n() / size() 上限／商の表の要素数
quotients() 異なる正の n/i を降順に並べた const vector<uint64_t>&
index(x) quotients() に含まれる正の x の添字
prime_sum_linear(a,b) 各要素が \(\displaystyle\sum_{p\leq q_i,\ p\text{ prime}}(a+bp)\) の vector<T> を返す
prime_sum(integer_prefix) 入力： \(\displaystyle\texttt{integer\_prefix}(x)=\sum_{k=1}^x g(k)\) を T で返す関数。\(g(1)=1\) かつ全ての正整数 \(a,b\) に対して \(g(ab)=g(a)g(b)\) を満たす完全乗法的な \(g\) に限る。返り値： 各要素が \(\displaystyle\sum_{p\leq q_i,\ p\text{ prime}}g(p)\) の vector<T>
sum(prime_prefix,prime_power) F(n) を T で返す
prefix_sum(prime_prefix,prime_power) 各要素が \(\displaystyle F(q_i)=\sum_{k=1}^{q_i}f(k)\) の vector<T> を返す
get(table,x) 作った表の x に対する値。x = 0 は 0
swap(other) 二つの計算領域を交換する

prime_power(p,e,pe) は \(f(p^e)\) を T で返します。p は素数、e >= 1、pe = p^e <= n です。同じ引数には同じ値を返してください。prime_prefix の素数上の値と、このコールバックの e = 1 の値は一致する必要があります。高い冪に対する値は任意に指定できます。

T は整数・modint など、0・1 の構築と加減乗算を備えた可換環の値を使います。コールバックの戻り値も T にしてください。逆元や除算は不要で、合成数 mod の型も使えます。prime_sum_linear では std::uint64_t から T への変換も使います。三角数は整数の片側を先に半分にするため、2 の逆元を要求しません。整数型の場合は、中間値もその型に収めてください。大きな整数和には __int128_t などを使えます。

代表的な指定は次の通りです。

関数 素数上の表 prime_power(p,e,pe)
常に 1 prime_sum_linear(T(1),T(0)) T(1)
恒等関数 prime_sum_linear(T(0),T(1)) T(pe)
トーティエント φ prime_sum_linear(T(-1),T(1)) T(pe-pe/p)
メビウス μ prime_sum_linear(T(-1),T(0)) e == 1 ? T(-1) : T(0)
約数個数 τ prime_sum_linear(T(2),T(0)) T(e+1)
約数和 σ prime_sum_linear(T(1),T(1)) \(\displaystyle\sum_{j=0}^{e}T(p)^j\)
平方数なら 1 prime_sum_linear(T(0),T(0)) T(e%2 == 0)

prime_sum(integer_prefix) では、補助関数 \(g\) が完全乗法的、すなわち互いに素という制約なしに \(g(ab)=g(a)g(b)\)、かつ \(g(1)=1\) であることが必要です。通常の乗法性だけでは不十分で、例えばトーティエント関数の累積和を直接渡すことはできません。integer_prefix(x) に \(\displaystyle\sum_{i=1}^x g(i)\) を与えると、各点で \(\displaystyle\sum_{p\leq x,\ p\text{ prime}}g(p)\) を返します。g(i)=1,i,i^2 や Dirichlet 指標などが使えます。返された表の線形結合を作れば、素数上の値が多項式などで表せる関数にも対応できます。別の方法で作った正しい素数上の表を sum・prefix_sum に渡しても構いません。

範囲と計算量

n は非負の std::uint64_t です。全ての表は quotients() と同じ長さ・順序です。get・index の正の引数はこの列に含まれる値だけに限ります。特に 1〜整数平方根までは全て含まれます。prefix_sum は長さ n の配列を作る操作ではありません。任意の上限 m の和は、上限 m のオブジェクトを構築して sum で求められます。

上限 0 では表は空で、sum と get(table,0) は 0 です。上限 1 では sum は 1 です。呼び出しはオブジェクトや入力の表を変更しないため、別の関数や異なる素数冪の値で計算領域を再利用できます。コピーは独立し、ムーブ後の元は上限 0 の状態になります。必要な領域は入力に応じて増えるため、時間とメモリに収まる範囲で使ってください。

商の表と素数篩を共有し、素数上の表は奇数に圧縮した Min_25 篩で作ります。prefix_sum は逆順の素数から高い冪を復元し、商が同じになる範囲をまとめて更新します。

sum は Black algorithm の列挙を使います。上限が \(2^{32}\) 以上で、T::mod() が 32 bit の法を返し、3 の逆元が存在する modint では、prime_sum・prime_sum_linear・sum を素数の範囲ごとに分割する処理へ自動で切り替えます。sum では平方因子を持たない数の和と強力数上の補正を組み合わせます。この処理は Library Checker の高速な提出 を参考にしています。型への追加要件はなく、それ以外の型でも同じ API を使えます。

T の演算とコールバックを \(O(1)\) とすると、構築は \(\displaystyle O(\sqrt n\log\log n)\) 時間・\(\displaystyle O(\sqrt n)\) 空間です。prime_sum・prime_sum_linear・sum・prefix_sum はそれぞれ \(\displaystyle O\!\left(\frac{n^{3/4}}{\log n}\right)\) 時間、保持する表と作業領域は \(\displaystyle O(\sqrt n)\) です。表の参照は \(O(1)\)、コピーは \(\displaystyle O(\sqrt n)\) 時間・空間です。n <= 1 は定数時間で処理します。

依存ライブラリ

検証コード

コード

#ifndef MULTIPLICATIVE_PREFIX_SUM_HPP
#define MULTIPLICATIVE_PREFIX_SUM_HPP

#include "../components/floor_quotients.hpp"
#include "../components/multiplicative_sum_backend.hpp"
#include "prime_sieve.hpp"
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <utility>
#include <vector>

template<class T> class MultiplicativePrefixSum {
	FloorQuotients quotients_;
	std::vector<std::size_t> primes_;

	bool use_backend(T& inverse) const {
		if constexpr (requires { std::uint64_t(T::mod()); T(0) == T(0); }) {
			const auto modulus = std::uint64_t(T::mod());
			if (max_n() >= (UINT64_C(1) << 32) && modulus > 1 && modulus <= UINT32_MAX && modulus % 3) {
				inverse = T((modulus * (modulus % 3 == 1 ? 2 : 1) + 1) / 3);
				return T(3) * inverse == T(1);
			}
		}
		return false;
	}

	template<class R, class Prefix> std::vector<R> prime_sum_impl(Prefix&& integer_prefix) const {
		const std::uint64_t n = max_n(), root = quotients_.root();
		const std::size_t large = quotients_.large_size(), length = size(), half = (root + 1) / 2;
		const R two = n >= 2 ? integer_prefix(2) - integer_prefix(1) : R(0);
		std::vector<R> result, small;
		result.reserve(length), small.reserve(half);
		auto initial = [&](std::uint64_t x) { return x == 1 ? R(0) : integer_prefix(x) - two * integer_prefix(x / 2) - R(1); };
		for (std::size_t i = 0; i < large; ++i) result.push_back(initial(quotients()[i]));
		for (std::size_t i = 0; i < half; ++i) small.push_back(initial(2 * i + 1));
		for (std::uint64_t p : primes_) {
			if (p == 2) continue;
			const R weight = integer_prefix(p) - integer_prefix(p - 1), before = small[(p - 2) / 2];
			const auto end = std::min<std::uint64_t>(large, n / (p * p));
			const auto split = std::min<std::uint64_t>(end, large / p);
			std::uint64_t i = 1;
			for (; i <= split; ++i) result[i - 1] -= weight * (result[i * p - 1] - before);
			for (; i <= end; ++i) result[i - 1] -= weight * (small[(n / (i * p) - 1) / 2] - before);
			std::size_t x = half;
			for (std::uint64_t q = (root / p - 1) | 1; q >= p; q -= 2) {
				const R delta = weight * (small[(q - 1) / 2] - before);
				for (const auto stop = (q * p + 1) / 2; x >= stop; --x) small[x - 1] -= delta;
			}
		}
		for (auto& value : result) value += two;
		for (auto x = root; x; --x) result.push_back(small[(x - 1) / 2] + (x >= 2 ? two : R(0)));
		return result;
	}


  public:
	explicit MultiplicativePrefixSum(std::uint64_t n = 0) : quotients_(n), primes_(prime_sieve(quotients_.root() + 1)) {}
	MultiplicativePrefixSum(const MultiplicativePrefixSum&) = default;
	MultiplicativePrefixSum& operator=(const MultiplicativePrefixSum&) = default;
	MultiplicativePrefixSum(MultiplicativePrefixSum&& other) noexcept : MultiplicativePrefixSum() { swap(other); }
	MultiplicativePrefixSum& operator=(MultiplicativePrefixSum&& other) noexcept {
		if (this != &other) { MultiplicativePrefixSum empty; swap(empty); swap(other); }
		return *this;
	}
	void swap(MultiplicativePrefixSum& other) noexcept { quotients_.swap(other.quotients_); primes_.swap(other.primes_); }
	std::uint64_t max_n() const { return quotients_.max_n(); }
	std::size_t size() const { return quotients_.size(); }
	const std::vector<std::uint64_t>& quotients() const { return quotients_.values(); }
	std::size_t index(std::uint64_t x) const { return quotients_.index(x); }
	T get(const std::vector<T>& table, std::uint64_t x) const {
		assert(table.size() == size());
		return x ? table[index(x)] : T(0);
	}
	template<class Prefix> std::vector<T> prime_sum(Prefix integer_prefix) const {
		T inverse(0);
		if (use_backend(inverse)) return multiplicative_sum_detail::Backend<T>(quotients_, primes_).prime_sum(integer_prefix, inverse);
		return prime_sum_impl<T>(integer_prefix);
	}
	std::vector<T> prime_sum_linear(T constant, T linear) const {
		auto identity_prefix = [](std::uint64_t x) -> T {
			return x & 1 ? T(x) * (T(x / 2) + T(1)) : T(x / 2) * (T(x) + T(1));
		};
		T inverse(0);
		const bool split = use_backend(inverse);
		if constexpr (requires { linear == T(0); }) {
			if (linear == T(0)) {
				if (constant == T(0)) return std::vector<T>(size(), T(0));
				if (split) {
					auto prefix = [](std::uint64_t x) { return T(x); };
					auto result = multiplicative_sum_detail::Backend<T>(quotients_, primes_).prime_sum(prefix, inverse);
					for (auto& value : result) value *= constant;
					return result;
				}
				auto counts = prime_sum_impl<std::uint64_t>([](std::uint64_t x) { return x; });
				std::vector<T> result;
				result.reserve(size());
				for (auto value : counts) result.push_back(T(value) * constant);
				return result;
			}
			if (constant == T(0)) {
				auto result = prime_sum(identity_prefix);
				for (auto& value : result) value *= linear;
				return result;
			}
		}
		if (split) {
			using Weights = multiplicative_sum_detail::LinearWeights<T>;
			auto prefix = [&](std::uint64_t x) { return Weights(T(x), identity_prefix(x)); };
			auto sums = multiplicative_sum_detail::Backend<Weights>(quotients_, primes_).prime_sum(prefix, Weights(inverse));
			std::vector<T> result;
			result.reserve(size());
			for (const auto& value : sums) result.push_back(value.count * constant + value.weighted * linear);
			return result;
		}
		const std::uint64_t n = max_n(), root = quotients_.root();
		const std::size_t large = quotients_.large_size(), length = size(), half = (root + 1) / 2;
		std::vector<T> result, small;
		std::vector<std::uint64_t> count;
		std::vector<std::uint32_t> count_small;
		result.reserve(length), small.reserve(half);
		count.reserve(large), count_small.reserve(half);
		for (std::size_t i = 0; i < large; ++i) {
			const auto x = quotients()[i], c = x / 2 + (x & 1);
			count.push_back(std::uint64_t(c - 1));
			result.push_back(T(c) * T(c) - T(1));
		}
		for (std::size_t i = 0; i < half; ++i) {
			count_small.push_back(std::uint32_t(i));
			small.push_back(T(i + 1) * T(i + 1) - T(1));
		}
		for (std::uint64_t p : primes_) {
			if (p == 2) continue;
			const T weight = T(p), before = small[(p - 2) / 2];
			const std::uint64_t count_before = count_small[(p - 2) / 2];
			const auto end = std::min<std::uint64_t>(large, n / (p * p));
			const auto split = std::min<std::uint64_t>(end, large / p);
			std::uint64_t i = 1;
			for (; i <= split; ++i) {
				count[i - 1] -= count[i * p - 1] - count_before;
				result[i - 1] -= weight * (result[i * p - 1] - before);
			}
			for (; i <= end; ++i) {
				const auto at = (n / (i * p) - 1) / 2;
				count[i - 1] -= count_small[at] - count_before;
				result[i - 1] -= weight * (small[at] - before);
			}
			std::size_t x = half;
			for (std::uint64_t q = (root / p - 1) | 1; q >= p; q -= 2) {
				const std::uint64_t count_delta = count_small[(q - 1) / 2] - count_before;
				const T delta = weight * (small[(q - 1) / 2] - before);
				for (const auto stop = (q * p + 1) / 2; x >= stop; --x) {
					count_small[x - 1] -= count_delta;
					small[x - 1] -= delta;
				}
			}
		}
		for (std::size_t i = 0; i < large; ++i) result[i] = (T(count[i]) + T(1)) * constant + (result[i] + T(2)) * linear;
		for (auto x = root; x; --x) result.push_back((T(count_small[(x - 1) / 2]) + T(x >= 2)) * constant + (small[(x - 1) / 2] + T(x >= 2 ? 2 : 0)) * linear);
		return result;
	}
	template<class PrimePower> std::vector<T> prefix_sum(const std::vector<T>& prime_prefix, PrimePower prime_power) const {
		assert(prime_prefix.size() == size());
		std::vector<T> result = prime_prefix, weights;
		std::vector<std::uint64_t> powers;
		weights.reserve(64), powers.reserve(64);
		for (std::size_t j = primes_.size(); j--;) {
			const std::uint64_t p = primes_[j];
			const T before = prime_prefix[index(p)];
			weights.clear(), powers.clear();
			for (std::uint64_t power = p;; power *= p) {
				powers.push_back(power);
				weights.push_back(prime_power(p, static_cast<unsigned>(powers.size()), power));
				if (power > max_n() / p) break;
			}
			const T first = weights[0], second = weights[1];
			const auto length = size(), large = quotients_.large_size();
			const auto end = std::min<std::uint64_t>(large, max_n() / (p * p));
			std::size_t i = 0;
			if (powers.size() > 2) for (; i < end && quotients()[i] >= powers[2]; ++i) {
				const auto d0 = (i + 1) * p;
				T delta = first * (result[d0 <= large ? d0 - 1 : length - max_n() / d0] - before) + second;
				for (std::size_t e = 1; e + 1 < powers.size() && powers[e + 1] <= quotients()[i]; ++e) {
					const auto d = (i + 1) * powers[e];
					delta += weights[e] * (result[d <= large ? d - 1 : length - max_n() / d] - before) + weights[e + 1];
				}
				result[i] += delta;
			}
			const auto split = std::min<std::uint64_t>(end, large / p);
			for (; i < split; ++i)
				result[i] += first * (result[(i + 1) * p - 1] - before) + second;
			for (; i < end; ++i)
				result[i] += first * (result[length - max_n() / ((i + 1) * p)] - before) + second;
			std::uint64_t x = quotients_.root();
			while (x >= p * p) {
				const auto q = x / p, stop = q * p;
				T delta = first * (result[length - q] - before) + second;
				for (std::size_t e = 1; e + 1 < powers.size() && powers[e + 1] <= x; ++e)
					delta += weights[e] * (result[length - x / powers[e]] - before) + weights[e + 1];
				for (; x >= stop; --x) result[length - x] += delta;
			}
		}
		for (auto& value : result) value += T(1);
		return result;
	}
	template<class PrimePower> T sum(const std::vector<T>& prime_prefix, PrimePower prime_power) const {
		assert(prime_prefix.size() == size());
		if (!max_n()) return T(0);
		T inverse(0);
		if (use_backend(inverse)) return multiplicative_sum_detail::Backend<T>(quotients_, primes_).sum(prime_prefix, prime_power, inverse);
		std::vector<T> first, second, square_prefix{T(0)};
		first.reserve(primes_.size()), second.reserve(primes_.size()), square_prefix.reserve(primes_.size() + 1);
		for (std::uint64_t p : primes_) {
			first.push_back(prime_power(p, 1, p));
			second.push_back(prime_power(p, 2, p * p));
			square_prefix.push_back(square_prefix.back() + second.back());
		}
		T result = prime_prefix[0] + T(1);
		const auto length = size(), large = quotients_.large_size();
		auto dfs = [&](auto&& self, std::size_t i, unsigned e, std::uint64_t product, std::uint64_t power) -> T {
			const std::uint64_t p = primes_[i], limit = max_n() / product;
			T delta = e == 1 ? second[i] : prime_power(p, e + 1, power * p);
			if (p * p <= limit) delta += self(self, i, e + 1, product * p, power * p);
			T weight = e == 1 ? first[i] : e == 2 ? second[i] : prime_power(p, e, power);
			if constexpr (requires { weight == T(0); }) if (weight == T(0)) return delta;
			T tail = prime_prefix[product <= large ? product - 1 : length - limit] - prime_prefix[length - p];
			std::size_t j = i + 1;
			for (; j < primes_.size(); ++j) {
				const std::uint64_t q = primes_[j];
				if (__uint128_t(q) * q * q > limit) break;
				tail += self(self, j, 1, product * q, q);
			}
			const auto begin = j;
			for (; j < primes_.size(); ++j) {
				const std::uint64_t q = primes_[j];
				if (q * q > limit) break;
				tail += first[j] * (prime_prefix[product * q <= large ? product * q - 1 : length - limit / q] - prime_prefix[length - q]);
			}
			return delta + weight * (tail + square_prefix[j] - square_prefix[begin]);
		};
		for (std::size_t i = 0; i < primes_.size(); ++i) result += dfs(dfs, i, 1, primes_[i], primes_[i]);
		return result;
	}
};

#endif
Back to top page