整数 k 乗根

math/kth_root.hpp · GitHub で見る

共通ページ

#include "lib/math/kth_root.hpp"

auto a = kth_root(215, 3);                     // 5
auto b = kth_root(216, 3);                     // 6
auto c = kth_root(18446744073709551615ULL, 2);  // 4294967295

関数	内容・返り値	計算量
kth_root(n,k)	\(\lfloor n^{1/k}\rfloor\) を uint64_t で返す	\(O(1)\)（64 bit 固定）

n と k は uint64_t で、\(0\le n<2^{64}\)、\(k\ge1\) が条件です。負数と k == 0 は扱いません。n == 0 の答えは 0、k == 1 の答えは n です。k は 64 を超えても構いません。

平方根・4 乗根・8 乗根は sqrt、それ以外の 3〜15 乗根は浮動小数の近似から候補を求め、整数演算で確定します。平方根の比較には __uint128_t、それ以外の累乗にはオーバーフローを検出する 64 bit 乗算を使います。候補が隣接する整数に補正できない場合は二分探索します。16 乗根以降は固定の小さな閾値表とビット長による判定で求めます。完全 k 乗数の直前でも丸め誤差は残りません。GCC・Clang と Mac の Apple Clang に対応します。

64 bit 固定のため時間・追加空間は \(O(1)\) です。3〜15 乗根で二分探索まで必要になった場合の整数乗算回数は、\(n\ge2\) として最悪 \(O((\log n/k)\log k)\) です。

検証コード

• test/kth_root.test.cpp
• test/multiplicative_prefix_sum_totient.test.cpp

コード

#ifndef KTH_ROOT_HPP
#define KTH_ROOT_HPP

#include <array>
#include <algorithm>
#include <limits>
#include <bit>
#include <cassert>
#include <cmath>
#include <cstdint>

namespace kth_root_detail {
inline constexpr auto thresholds = [] {
	std::array<std::array<std::uint64_t, 13>, 25> table{};
	for (unsigned k = 16; k <= 40; ++k) for (unsigned a = 3; a <= 15; ++a) {
		std::uint64_t p = 1;
		for (unsigned i = 0; i < k; ++i) {
			if (p > UINT64_MAX / a) { p = 0; break; }
			p *= a;
		}
		table[k - 16][a - 3] = p - 1;
	}
	return table;
}();

inline std::uint64_t cube(std::uint64_t n) {
	if constexpr (std::numeric_limits<double>::is_iec559 && std::numeric_limits<double>::digits == 53 && sizeof(double) == sizeof(std::uint64_t)) {
		static constexpr auto inverse_roots = [] {
			std::array<double, 64> result{};
			for (int i = 0; i < 64; ++i) {
				double a = 1 + (i + 0.5) / 64, x = 1;
				for (int j = 0; j < 8; ++j) x += x * (1 - a * x * x * x) / 3;
				result[i] = x;
			}
			return result;
		}();
		double a = double(n);
		std::uint64_t bits = std::bit_cast<std::uint64_t>(a);
		unsigned exponent = unsigned(bits >> 52) - 1023;
		constexpr double factor[] = {1, 0.79370052598409973738, 0.62996052494743658238};
		double x = inverse_roots[(bits >> 46) & 63] * factor[exponent % 3]
			* std::bit_cast<double>(std::uint64_t(1023 - exponent / 3) << 52);
		{
			double h = 1 - a * x * x * x;
			x += x * h * (1.0 / 3 + h * (2.0 / 9 + h * (14.0 / 81)));
		}
		std::uint64_t r = std::min(std::uint64_t(1 / x), std::uint64_t(2642245));
		while (r * r * r > n) --r;
		while (r < 2642245 && (r + 1) * (r + 1) * (r + 1) <= n) ++r;
		return r;
	} else {
		std::uint64_t r = std::cbrt(double(n));
		while (__uint128_t(r) * r * r > n) --r;
		while (__uint128_t(r + 1) * (r + 1) * (r + 1) <= n) ++r;
		return r;
	}
}

template<unsigned K>
constexpr double power(double x) {
	double value = 1;
	for (unsigned k = K;; x *= x) {
		if (k & 1) value *= x;
		if (!(k >>= 1)) return value;
	}
}

template<unsigned K>
inline double estimate(std::uint64_t n) {
	if constexpr (std::numeric_limits<double>::is_iec559 && std::numeric_limits<double>::digits == 53 && sizeof(double) == sizeof(std::uint64_t)) {
		static constexpr auto roots = [] {
			std::array<double, 64> values{};
			for (int i = 0; i < 64; ++i) {
				double a = 1 + (i + 0.5) / 64, x = 1;
				for (int j = 0; j < 10; ++j) x += x * (1 - a * power<K>(x)) / K;
				values[i] = x;
			}
			return values;
		}();
		static constexpr auto factors = [] {
			std::array<double, K> values{};
			double x = 1;
			for (int j = 0; j < 10; ++j) x += x * (1 - 2 * power<K>(x)) / K;
			values[0] = 1;
			for (unsigned j = 1; j < K; ++j) values[j] = values[j - 1] * x;
			return values;
		}();
		double a = double(n);
		auto bits = std::bit_cast<std::uint64_t>(a);
		unsigned exponent = unsigned(bits >> 52) - 1023;
		double x = roots[(bits >> 46) & 63] * factors[exponent % K]
			* std::bit_cast<double>(std::uint64_t(1023 - exponent / K) << 52);
		double h = 1 - a * power<K>(x);
		constexpr double c = 1.0 / K;
		x += x * h * (c + h * (c * (c + 1) / 2 + h * (c * (c + 1) * (c + 2) / 6)));
		return 1 / x;
	} else return std::pow(double(n), 1.0 / K);
}

template <unsigned K>
inline std::uint64_t small(std::uint64_t n) {
	auto leq = [n](std::uint64_t x) {
		std::uint64_t a = x, r = 1;
		for (unsigned k = K;;) {
			if ((k & 1) && __builtin_mul_overflow(r, a, &r)) return false;
			if (!(k >>= 1)) return r <= n;
			if (__builtin_mul_overflow(a, a, &a)) return false;
		}
	};
	std::uint64_t lo = 1, hi = std::uint64_t(1) << ((std::bit_width(n) - 1) / K + 1);
	double estimate;
	if constexpr (K == 4) estimate = std::sqrt(std::sqrt(double(n)));
	else if constexpr (K == 8) estimate = std::sqrt(std::sqrt(std::sqrt(double(n))));
	else estimate = kth_root_detail::estimate<K>(n);
	std::uint64_t r = std::uint64_t(estimate);
	if (r >= hi) r = hi - 1;
	if (r < lo) r = lo;
	if (leq(r)) {
		if (!leq(r + 1)) return r;
		lo = r + 1;
	} else {
		if (leq(r - 1)) return r - 1;
		hi = r - 1;
	}
	while (lo + 1 < hi) {
		std::uint64_t mid = (lo + hi) / 2;
		if (leq(mid)) lo = mid;
		else hi = mid;
	}
	return lo;
}
}

inline std::uint64_t kth_root(std::uint64_t n, std::uint64_t k) {
	assert(k > 0);
	if (k == 1 || n <= 1) return n;
	if (k >= unsigned(std::bit_width(n))) return 1;
	if (k == 2) {
		std::uint64_t r = std::sqrt(double(n));
		if (__uint128_t(r) * r > n) --r;
		if (__uint128_t(r + 1) * (r + 1) <= n) ++r;
		return r;
	}
	if (k >= 16) {
		if (k > 40) return 2;
		unsigned lo = 3, hi = 1U << (63 / k + 1);
		while (lo < hi) {
			unsigned mid = (lo + hi) / 2;
			if (n > kth_root_detail::thresholds[k - 16][mid - 3]) lo = mid + 1;
			else hi = mid;
		}
		return lo - 1;
	}
	switch (k) {
	case 3: return kth_root_detail::cube(n);
	case 4: return kth_root_detail::small<4>(n);
	case 5: return kth_root_detail::small<5>(n);
	case 6: return kth_root_detail::small<6>(n);
	case 7: return kth_root_detail::small<7>(n);
	case 8: return kth_root_detail::small<8>(n);
	case 9: return kth_root_detail::small<9>(n);
	case 10: return kth_root_detail::small<10>(n);
	case 11: return kth_root_detail::small<11>(n);
	case 12: return kth_root_detail::small<12>(n);
	case 13: return kth_root_detail::small<13>(n);
	case 14: return kth_root_detail::small<14>(n);
	default: return kth_root_detail::small<15>(n);
	}
}

#endif

Back to top page