haar_lib/math/
kth_root.rs

1//! Kth root
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/kth_root_integer>
5
6/// $\lfloor a^{1/k} \rfloor$を求める。
7pub fn kth_root(a: u64, k: u64) -> u64 {
8    assert!(k >= 1);
9
10    match (a, k) {
11        (0, _) => 0,
12        (1, _) => 1,
13        (a, 1) => a,
14        (_, k) if k > 64 => 1,
15        (a, k) => {
16            let mut lb = 0;
17            let mut ub = a;
18            while ub - lb > 1 {
19                let mid = midpoint(ub, lb);
20                if check(mid, k, a) {
21                    lb = mid;
22                } else {
23                    ub = mid;
24                }
25            }
26
27            lb
28        }
29    }
30}
31
32fn check(mut a: u64, mut k: u64, n: u64) -> bool {
33    let mut r: u64 = 1;
34
35    while k > 0 {
36        if k % 2 == 1 {
37            let Some(_r) = r.checked_mul(a) else {
38                return false;
39            };
40            r = _r;
41        }
42        if let Some(_a) = a.checked_mul(a) {
43            a = _a;
44        } else if k > 1 {
45            return false;
46        }
47        k >>= 1;
48    }
49
50    r <= n
51}
52
53fn midpoint(a: u64, b: u64) -> u64 {
54    ((a ^ b) >> 1) + (a & b)
55}
haar_lib/math/kth_root.rs

haar_lib/math/
kth_root.rs
