:heavy_check_mark: k乗根(整数) (Math/kth_root_integer.hpp)

kth_root

T kth_root(T a, T k)

入力された値の$\operatorname{floor}\left( \sqrt[k]{a} \right)$を返します。

制約

  • Tは整数型
  • $0 \leq a$
  • $0 \lt k$

計算量

  • $O(k \log a)$

Depends on

Verified with

Code

#include <bits/stdc++.h>
using namespace std;

#include "Other/binary_search.hpp"

template <typename T>
enable_if_t<is_integral_v<T>, T>
kth_root(T a, T k)
{
    assert(0 <= a);
    assert(0 < k);

    if (a == 1)
    {
        return 1;
    }
    if (k == 1)
    {
        return a;
    }

    auto check_root = [&](T x) -> bool
    {
        T tmp = 1;
        for (T i = 0; i < k; i++)
        {
            if (tmp > a / x)
            {
                return false;
            }

            tmp *= x;
        }

        return tmp <= a;
    };

    return bin_search<T>(0, a, check_root);
}
#line 1 "Math/kth_root_integer.hpp"
#include <bits/stdc++.h>
using namespace std;

#line 2 "Other/binary_search.hpp"
using namespace std;

template <typename T>
enable_if_t<is_integral_v<T>, T>
bin_search(T ok, T ng, function<bool(T)> check)
{
    while (max(ok, ng) - min(ok, ng) > 1)
    {
        T mid = midpoint(ok, ng);
        (check(mid) ? ok : ng) = mid;
    }
    return ok;
}

template <typename T>
enable_if_t<is_floating_point_v<T>, T>
bin_search(T ok, T ng, function<bool(T)> check)
{
    // 大きすぎるとcheck(x)のxが不当に大きくなるのでよき大きさに切る。
    assert(abs(ok) < 1e20 && abs(ng) < 1e20);

    for (int i = 0; i < 100; i++)
    {
        T mid = midpoint(ok, ng);
        (check(mid) ? ok : ng) = mid;
    }
    return ok;
}
#line 5 "Math/kth_root_integer.hpp"

template <typename T>
enable_if_t<is_integral_v<T>, T>
kth_root(T a, T k)
{
    assert(0 <= a);
    assert(0 < k);

    if (a == 1)
    {
        return 1;
    }
    if (k == 1)
    {
        return a;
    }

    auto check_root = [&](T x) -> bool
    {
        T tmp = 1;
        for (T i = 0; i < k; i++)
        {
            if (tmp > a / x)
            {
                return false;
            }

            tmp *= x;
        }

        return tmp <= a;
    };

    return bin_search<T>(0, a, check_root);
}
Back to top page