Kth root integer
(Mylib/Number/kth_root_integer.cpp)

• View this file on GitHub
• Last update: 2021-04-23 23:44:44+09:00
• Link: View error logs on GitHub Actions

Operations

Requirements

Notes

Problems

References

Verified with

• test/yosupo-judge/kth_root_integer/main.test.cpp

Code

#pragma once
#include <cassert>
#include <cstdint>

namespace haar_lib {
  uint64_t kth_root(uint64_t a, int k) {
    assert(k >= 1);
    if (k == 1) return a;
    if (a == 1) return 1;

    uint64_t lb = 0, ub = a;

    auto check =
        [](uint64_t a, int k, uint64_t n) {
          uint64_t r = 1;

          while (k > 0) {
            if (k & 1) {
              if (__builtin_umulll_overflow(r, a, (unsigned long long int*) &r)) return false;
            }
            if (__builtin_umulll_overflow(a, a, (unsigned long long int*) &a) and k > 1) return false;
            k >>= 1;
          }

          return r <= n;
        };

    while (ub - lb > 1) {
      uint64_t mid = lb + (ub - lb) / 2;

      if (check(mid, k, a)) {
        lb = mid;
      } else {
        ub = mid;
      }
    }

    return lb;
  }
}  // namespace haar_lib

#line 2 "Mylib/Number/kth_root_integer.cpp"
#include <cassert>
#include <cstdint>

namespace haar_lib {
  uint64_t kth_root(uint64_t a, int k) {
    assert(k >= 1);
    if (k == 1) return a;
    if (a == 1) return 1;

    uint64_t lb = 0, ub = a;

    auto check =
        [](uint64_t a, int k, uint64_t n) {
          uint64_t r = 1;

          while (k > 0) {
            if (k & 1) {
              if (__builtin_umulll_overflow(r, a, (unsigned long long int*) &r)) return false;
            }
            if (__builtin_umulll_overflow(a, a, (unsigned long long int*) &a) and k > 1) return false;
            k >>= 1;
          }

          return r <= n;
        };

    while (ub - lb > 1) {
      uint64_t mid = lb + (ub - lb) / 2;

      if (check(mid, k, a)) {
        lb = mid;
      } else {
        ub = mid;
      }
    }

    return lb;
  }
}  // namespace haar_lib

Back to top page