tests/LibraryChecker_Kth_root.test.cpp

View this file on GitHub
Last update: 2026-11-04 14:46:41+09:00
Problem: https://judge.yosupo.jp/problem/kth_root_integer

Depends on

kth_root (Others/kth_root.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/kth_root_integer"

#include "kth_root.hpp"

int main() {
    cin.tie(nullptr);
    ios_base::sync_with_stdio(false);
    
    int T;
    cin >> T;
    while (T--) {
        uint64_t A, K;
        cin >> A >> K;
        cout << kth_root(A, K) << '\n';
    }
}

#line 1 "tests/LibraryChecker_Kth_root.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/kth_root_integer"

#line 2 "Others/kth_root.hpp"
#include <bits/stdc++.h>
#line 4 "Others/kth_root.hpp"
using namespace std;

// kth root
// N < 2^64, K <= 64
uint64_t kth_root(uint64_t N, uint64_t K = 2) {
    assert(K >= 1);
    if (N <= 1 || K == 1) return N;
    if (K >= 64) return 1;
    if (N == uint64_t(-1)) --N;
    
    auto mul = [&](uint64_t x, uint64_t y) -> uint64_t {
        if (x < UINT_MAX && y < UINT_MAX) return x * y;
        if (x == uint64_t(-1) || y == uint64_t(-1)) return uint64_t(-1);
        return (x <= uint64_t(-1) / y ? x * y : uint64_t(-1));
    };
    auto power = [&](uint64_t x, uint64_t k) -> uint64_t {
        if (k == 0) return 1ULL;
        uint64_t res = 1ULL;
        while (k) {
            if (k & 1) res = mul(res, x);
            x = mul(x, x);
            k >>= 1;
        }
        return res;
    };
    
    uint64_t res;
    if (K == 2) res = sqrtl(N) - 1;
    else if (K == 3) res = cbrt(N) - 1;
    else res = pow(N, nextafter(1 / double(K), 0));
    while (power(res + 1, K) <= N) ++res;
    return res;
}
#line 4 "tests/LibraryChecker_Kth_root.test.cpp"

int main() {
    cin.tie(nullptr);
    ios_base::sync_with_stdio(false);
    
    int T;
    cin >> T;
    while (T--) {
        uint64_t A, K;
        cin >> A >> K;
        cout << kth_root(A, K) << '\n';
    }
}