tests/floor_kth_root.test.cpp

• View this file on GitHub
• Last update: 2025-01-22 03:51:11+00:00
• Problem: https://judge.yosupo.jp/problem/kth_root_integer

Depends on

• $\left\lfloor x^\frac{1}{k} \right\rfloor$ (tools/floor_kth_root.hpp)
• $\left\lfloor \sqrt{x} \right\rfloor$ (tools/floor_sqrt.hpp)
• std::gcd(m, n) extended for my library (tools/gcd.hpp)
• Check whether T is a monoid (tools/is_monoid.hpp)
• Typical monoids (tools/monoid.hpp)
• $b^n$ under a given monoid, and std::pow(b, n) extended for my library (tools/pow.hpp)
• $\mathbb{Z} \cup \{\infty, -\infty, \mathrm{NaN}\}$ and $\mathbb{Z}_{\geq 0} \cup \{\infty, \mathrm{NaN}\}$ (tools/safe_int.hpp)
• $x^2$ under a given monoid (tools/square.hpp)

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/kth_root_integer

#include <iostream>
#include "tools/floor_kth_root.hpp"

using ull = unsigned long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int T;
  std::cin >> T;
  for (int t = 0; t < T; ++t) {
    ull A;
    int K;
    std::cin >> A >> K;
    std::cout << tools::floor_kth_root(A, K) << '\n';
  }

  return 0;
}

#line 1 "tests/floor_kth_root.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/kth_root_integer

#include <iostream>
#line 1 "tools/floor_kth_root.hpp"



#include <cassert>
#line 1 "tools/floor_sqrt.hpp"



#line 5 "tools/floor_sqrt.hpp"

namespace tools {

  template <typename T>
  T floor_sqrt(const T n) {
    assert(n >= 0);

    T ok = 0;
    T ng;
    for (ng = 1; ng <= n / ng; ng *= 2);

    while (ng - ok > 1) {
      const T mid = ok + (ng - ok) / 2;
      if (mid <= n / mid) {
        ok = mid;
      } else {
        ng = mid;
      }
    }

    return ok;
  }
}


#line 1 "tools/pow.hpp"



#include <type_traits>
#line 6 "tools/pow.hpp"
#include <cmath>
#line 1 "tools/monoid.hpp"



#line 5 "tools/monoid.hpp"
#include <algorithm>
#include <limits>
#line 1 "tools/gcd.hpp"



#line 5 "tools/gcd.hpp"
#include <numeric>

namespace tools {
  template <typename M, typename N>
  constexpr ::std::common_type_t<M, N> gcd(const M m, const N n) {
    return ::std::gcd(m, n);
  }
}


#line 9 "tools/monoid.hpp"

namespace tools {
  namespace monoid {
    template <typename M, M ...dummy>
    struct max;

    template <typename M>
    struct max<M> {
      static_assert(::std::is_arithmetic_v<M>, "M must be a built-in arithmetic type.");

      using T = M;
      static T op(const T lhs, const T rhs) {
        return ::std::max(lhs, rhs);
      }
      static T e() {
        if constexpr (::std::is_integral_v<M>) {
          return ::std::numeric_limits<M>::min();
        } else {
          return -::std::numeric_limits<M>::infinity();
        }
      }
    };

    template <typename M, M E>
    struct max<M, E> {
      static_assert(::std::is_integral_v<M>, "M must be a built-in integral type.");

      using T = M;
      static T op(const T lhs, const T rhs) {
        assert(E <= lhs);
        assert(E <= rhs);
        return ::std::max(lhs, rhs);
      }
      static T e() {
        return E;
      }
    };

    template <typename M, M ...dummy>
    struct min;

    template <typename M>
    struct min<M> {
      static_assert(::std::is_arithmetic_v<M>, "M must be a built-in arithmetic type.");

      using T = M;
      static T op(const T lhs, const T rhs) {
        return ::std::min(lhs, rhs);
      }
      static T e() {
        if constexpr (::std::is_integral_v<M>) {
          return ::std::numeric_limits<M>::max();
        } else {
          return ::std::numeric_limits<M>::infinity();
        }
      }
    };

    template <typename M, M E>
    struct min<M, E> {
      static_assert(::std::is_integral_v<M>, "M must be a built-in integral type.");

      using T = M;
      static T op(const T lhs, const T rhs) {
        assert(lhs <= E);
        assert(rhs <= E);
        return ::std::min(lhs, rhs);
      }
      static T e() {
        return E;
      }
    };

    template <typename M>
    struct multiplies {
    private:
      using VR = ::std::conditional_t<::std::is_arithmetic_v<M>, const M, const M&>;

    public:
      using T = M;
      static T op(VR lhs, VR rhs) {
        return lhs * rhs;
      }
      static T e() {
        return T(1);
      }
    };

    template <>
    struct multiplies<bool> {
      using T = bool;
      static T op(const bool lhs, const bool rhs) {
        return lhs && rhs;
      }
      static T e() {
        return true;
      }
    };

    template <typename M>
    struct gcd {
    private:
      static_assert(!::std::is_arithmetic_v<M> || (::std::is_integral_v<M> && !::std::is_same_v<M, bool>), "If M is a built-in arithmetic type, it must be integral except for bool.");
      using VR = ::std::conditional_t<::std::is_arithmetic_v<M>, const M, const M&>;

    public:
      using T = M;
      static T op(VR lhs, VR rhs) {
        return ::tools::gcd(lhs, rhs);
      }
      static T e() {
        return T(0);
      }
    };

    template <typename M, M E>
    struct update {
      static_assert(::std::is_integral_v<M>, "M must be a built-in integral type.");

      using T = M;
      static T op(const T lhs, const T rhs) {
        return lhs == E ? rhs : lhs;
      }
      static T e() {
        return E;
      }
    };
  }
}


#line 1 "tools/square.hpp"



#line 1 "tools/is_monoid.hpp"



#line 5 "tools/is_monoid.hpp"
#include <utility>

namespace tools {

  template <typename M, typename = void>
  struct is_monoid : ::std::false_type {};

  template <typename M>
  struct is_monoid<M, ::std::enable_if_t<
    ::std::is_same_v<typename M::T, decltype(M::op(::std::declval<typename M::T>(), ::std::declval<typename M::T>()))> &&
    ::std::is_same_v<typename M::T, decltype(M::e())>
  , void>> : ::std::true_type {};

  template <typename M>
  inline constexpr bool is_monoid_v = ::tools::is_monoid<M>::value;
}


#line 6 "tools/square.hpp"

namespace tools {

  template <typename M>
  ::std::enable_if_t<::tools::is_monoid_v<M>, typename M::T> square(const typename M::T& x) {
    return M::op(x, x);
  }

  template <typename T>
  ::std::enable_if_t<!::tools::is_monoid_v<T>, T> square(const T& x) {
    return x * x;
  }
}


#line 9 "tools/pow.hpp"

namespace tools {

  template <typename M, typename E>
  ::std::enable_if_t<::std::is_integral_v<E>, typename M::T> pow(const typename M::T& base, const E exponent) {
    assert(exponent >= 0);
    return exponent == 0
      ? M::e()
      : exponent % 2 == 0
        ? ::tools::square<M>(::tools::pow<M>(base, exponent / 2))
        : M::op(::tools::pow<M>(base, exponent - 1), base);
  }

  template <typename T, typename E>
  ::std::enable_if_t<::std::is_integral_v<E>, T> pow(const T& base, const E exponent) {
    assert(exponent >= 0);
    return ::tools::pow<::tools::monoid::multiplies<T>>(base, exponent);
  }

  template <typename T, typename E>
  auto pow(const T base, const E exponent) -> ::std::enable_if_t<!::std::is_integral_v<E>, decltype(::std::pow(base, exponent))> {
    return ::std::pow(base, exponent);
  }
}


#line 1 "tools/safe_int.hpp"



#line 5 "tools/safe_int.hpp"
#include <cstddef>
#line 8 "tools/safe_int.hpp"
#include <array>
#include <optional>
#line 11 "tools/safe_int.hpp"

namespace tools {
  template <typename T, typename = void>
  class safe_int;

  template <typename T>
  class safe_int<T, ::std::enable_if_t<::std::is_signed_v<T>>> {
  private:
    enum class type {
      finite,
      pos_inf,
      neg_inf,
      nan
    };
    typename ::tools::safe_int<T>::type m_type;
    T m_value;

    constexpr safe_int(const typename ::tools::safe_int<T>::type type) :
      m_type(type), m_value(T()) {
    }

  public:
    constexpr safe_int() :
      m_type(::tools::safe_int<T>::type::finite), m_value(T()) {
    }
    explicit constexpr safe_int(const T value) :
      m_type(::tools::safe_int<T>::type::finite), m_value(value) {
    }
    constexpr safe_int(const ::tools::safe_int<T>& other) :
      m_type(other.m_type), m_value(other.m_value) {
    }
    ~safe_int() = default;
    constexpr ::tools::safe_int<T>& operator=(const ::tools::safe_int<T>& other) {
      this->m_type = other.m_type;
      this->m_value = other.m_value;
      return *this;
    }

    static constexpr ::tools::safe_int<T> infinity() {
      return tools::safe_int<T>(::tools::safe_int<T>::type::pos_inf);
    }
    static constexpr ::tools::safe_int<T> nan() {
      return tools::safe_int<T>(::tools::safe_int<T>::type::nan);
    }

  private:
    static constexpr int f1(const ::tools::safe_int<T>& n) {
      switch (n.m_type) {
      case ::tools::safe_int<T>::type::neg_inf:
        return 0;
      case ::tools::safe_int<T>::type::finite:
        return 1;
      case ::tools::safe_int<T>::type::pos_inf:
        return 2;
      default: // nan
        return 3;
      }
    };
    static constexpr int f2(const ::tools::safe_int<T>& n) {
      switch (n.m_type) {
      case ::tools::safe_int<T>::type::neg_inf:
        return 0;
      case ::tools::safe_int<T>::type::finite:
        if (n.m_value < 0) {
          return 1;
        } else if (n.m_value == 0) {
          return 2;
        } else {
          return 3;
        }
      case ::tools::safe_int<T>::type::pos_inf:
        return 4;
      default: // nan
        return 5;
      }
    };
    static constexpr ::std::optional<::tools::safe_int<T>> Q() {
      return ::std::nullopt;
    }
    static constexpr ::std::optional<::tools::safe_int<T>> Z() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(0));
    }
    static constexpr ::std::optional<::tools::safe_int<T>> N() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(::tools::safe_int<T>::type::neg_inf));
    }
    static constexpr ::std::optional<::tools::safe_int<T>> P() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(::tools::safe_int<T>::type::pos_inf));
    }
    static constexpr ::std::optional<::tools::safe_int<T>> U() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(::tools::safe_int<T>::type::nan));
    }
    static constexpr ::std::optional<bool> BQ() {
      return ::std::nullopt;
    }
    static constexpr ::std::optional<bool> BF() {
      return ::std::optional<bool>(false);
    }
    static constexpr ::std::optional<bool> BT() {
      return ::std::optional<bool>(true);
    }

  public:
    constexpr bool is_finite() const {
      return this->m_type == ::tools::safe_int<T>::type::finite;
    }

    constexpr bool is_nan() const {
      return this->m_type == ::tools::safe_int<T>::type::nan;
    }

    constexpr T val() const {
      assert(this->is_finite());
      return this->m_value;
    }

    friend constexpr bool operator==(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<bool>, 4>, 4>({{
        {BT(), BF(), BF(), BF()},
        {BF(), BQ(), BF(), BF()},
        {BF(), BF(), BT(), BF()},
        {BF(), BF(), BF(), BF()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      return x.m_value == y.m_value;
    }
    friend constexpr bool operator!=(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      return !(x == y);
    }

    constexpr ::tools::safe_int<T> operator+() const {
      return *this;
    }
    constexpr ::tools::safe_int<T> operator-() const {
      constexpr auto table = ::std::array<::std::optional<::tools::safe_int<T>>, 4>({
        {P(), Q(), N(), U()}
      });
      if (const auto r = table[f1(*this)]; r) return *r;

      if (this->m_value == ::std::numeric_limits<T>::min()) return *U();
      return ::tools::safe_int<T>(-this->m_value);
    }

    friend constexpr ::tools::safe_int<T> operator+(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 4>, 4>({{
        {N(), N(), U(), U()},
        {N(), Q(), P(), U()},
        {U(), P(), P(), U()},
        {U(), U(), U(), U()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      if (y.m_value > 0 && x.m_value > ::std::numeric_limits<T>::max() - y.m_value) return *U();
      if (y.m_value < 0 && x.m_value < ::std::numeric_limits<T>::min() - y.m_value) return *U();
      return ::tools::safe_int<T>(x.m_value + y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator+(const ::tools::safe_int<T>& x, const T& y) {
      return x + tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator+(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) + y;
    }

    friend constexpr ::tools::safe_int<T> operator-(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 4>, 4>({{
        {U(), N(), N(), U()},
        {P(), Q(), N(), U()},
        {P(), P(), U(), U()},
        {U(), U(), U(), U()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      if (y.m_value < 0 && x.m_value > ::std::numeric_limits<T>::max() + y.m_value) return *U();
      if (y.m_value > 0 && x.m_value < ::std::numeric_limits<T>::min() + y.m_value) return *U();
      return ::tools::safe_int<T>(x.m_value - y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator-(const ::tools::safe_int<T>& x, const T& y) {
      return x - tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator-(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) - y;
    }

    friend constexpr ::tools::safe_int<T> operator*(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 6>, 6>({{
        {P(), P(), U(), N(), N(), U()},
        {P(), Q(), Z(), Q(), N(), U()},
        {U(), Z(), Z(), Z(), U(), U()},
        {N(), Q(), Z(), Q(), P(), U()},
        {N(), N(), U(), P(), P(), U()},
        {U(), U(), U(), U(), U(), U()}
      }});
      if (const auto r = table[f2(x)][f2(y)]; r) return *r;

      if (x.m_value > 0) {
        if (y.m_value > 0) {
          if (x.m_value > ::std::numeric_limits<T>::max() / y.m_value) {
            return *U();
          }
        } else {
          if (y.m_value < ::std::numeric_limits<T>::min() / x.m_value) {
            return *U();
          }
        }
      } else {
        if (y.m_value > 0) {
          if (x.m_value < ::std::numeric_limits<T>::min() / y.m_value) {
            return *U();
          }
        } else {
          if (x.m_value != 0 && y.m_value < ::std::numeric_limits<T>::max() / x.m_value) {
            return *U();
          }
        }
      }
      return ::tools::safe_int<T>(x.m_value * y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator*(const ::tools::safe_int<T>& x, const T& y) {
      return x * tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator*(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) * y;
    }

    friend constexpr ::tools::safe_int<T> operator/(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 6>, 6>({{
        {U(), P(), U(), N(), U(), U()},
        {Z(), Q(), U(), Q(), Z(), U()},
        {Z(), Z(), U(), Z(), Z(), U()},
        {Z(), Q(), U(), Q(), Z(), U()},
        {U(), N(), U(), P(), U(), U()},
        {U(), U(), U(), U(), U(), U()}
      }});
      if (const auto r = table[f2(x)][f2(y)]; r) return *r;

      if (x.m_value == ::std::numeric_limits<T>::min() && y.m_value == -1) return *U();
      return ::tools::safe_int<T>(x.m_value / y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator/(const ::tools::safe_int<T>& x, const T& y) {
      return x / tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator/(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) / y;
    }

    friend constexpr ::tools::safe_int<T> operator%(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 6>, 6>({{
        {U(), U(), U(), U(), U(), U()},
        {U(), Q(), U(), Q(), U(), U()},
        {U(), Z(), U(), Z(), U(), U()},
        {U(), Q(), U(), Q(), U(), U()},
        {U(), U(), U(), U(), U(), U()},
        {U(), U(), U(), U(), U(), U()}
      }});
      if (const auto r = table[f2(x)][f2(y)]; r) return *r;

      if (x.m_value == ::std::numeric_limits<T>::min() && y.m_value == -1) return *U();
      return ::tools::safe_int<T>(x.m_value % y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator%(const ::tools::safe_int<T>& x, const T& y) {
      return x % tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator%(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) % y;
    }

    constexpr ::tools::safe_int<T>& operator+=(const ::tools::safe_int<T>& other) {
      return *this = *this + other;
    }
    constexpr ::tools::safe_int<T>& operator+=(const T& other) {
      return *this = *this + ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator-=(const ::tools::safe_int<T>& other) {
      return *this = *this - other;
    }
    constexpr ::tools::safe_int<T>& operator-=(const T& other) {
      return *this = *this - ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator*=(const ::tools::safe_int<T>& other) {
      return *this = *this * other;
    }
    constexpr ::tools::safe_int<T>& operator*=(const T& other) {
      return *this = *this * ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator/=(const ::tools::safe_int<T>& other) {
      return *this = *this / other;
    }
    constexpr ::tools::safe_int<T>& operator/=(const T& other) {
      return *this = *this / ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator%=(const ::tools::safe_int<T>& other) {
      return *this = *this % other;
    }
    constexpr ::tools::safe_int<T>& operator%=(const T& other) {
      return *this = *this % ::tools::safe_int<T>(other);
    }

    constexpr ::tools::safe_int<T>& operator++() {
      return *this += ::tools::safe_int<T>(T(1));
    }
    constexpr ::tools::safe_int<T> operator++(int) {
      const auto r = *this;
      ++(*this);
      return r;
    }
    constexpr ::tools::safe_int<T>& operator--() {
      return *this -= ::tools::safe_int<T>(T(1));
    }
    constexpr ::tools::safe_int<T> operator--(int) {
      const auto r = *this;
      --(*this);
      return r;
    }

    friend constexpr bool operator<(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<bool>, 4>, 4>({{
        {BF(), BT(), BT(), BF()},
        {BF(), BQ(), BT(), BF()},
        {BF(), BF(), BF(), BF()},
        {BF(), BF(), BF(), BF()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      return x.m_value < y.m_value;
    }
    friend constexpr bool operator>(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<bool>, 4>, 4>({{
        {BF(), BF(), BF(), BF()},
        {BT(), BQ(), BF(), BF()},
        {BT(), BT(), BF(), BF()},
        {BF(), BF(), BF(), BF()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      return x.m_value > y.m_value;
    }
    friend constexpr bool operator<=(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      return x < y || x == y;
    }
    friend constexpr bool operator>=(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      return x > y || x == y;
    }

    friend ::std::istream& operator>>(::std::istream& is, ::tools::safe_int<T>& self) {
      self.m_type = ::tools::safe_int<T>::type::finite;
      return is >> self.m_value;
    }
    friend ::std::ostream& operator<<(::std::ostream& os, const ::tools::safe_int<T>& self) {
      switch (self.m_type) {
      case ::tools::safe_int<T>::type::neg_inf:
        return os << "-inf";
      case ::tools::safe_int<T>::type::finite:
        return os << self.m_value;
      case ::tools::safe_int<T>::type::pos_inf:
        return os << "inf";
      default: // nan
        return os << "nan";
      }
    }
  };

  template <typename T>
  class safe_int<T, ::std::enable_if_t<::std::is_unsigned_v<T>>> {
  private:
    enum class type {
      finite,
      pos_inf,
      nan
    };
    typename ::tools::safe_int<T>::type m_type;
    T m_value;

    constexpr safe_int(const typename ::tools::safe_int<T>::type type) :
      m_type(type), m_value(T()) {
    }

  public:
    constexpr safe_int() :
      m_type(::tools::safe_int<T>::type::finite), m_value(T()) {
    }
    explicit constexpr safe_int(const T value) :
      m_type(::tools::safe_int<T>::type::finite), m_value(value) {
    }
    constexpr safe_int(const ::tools::safe_int<T>& other) :
      m_type(other.m_type), m_value(other.m_value) {
    }
    ~safe_int() = default;
    constexpr ::tools::safe_int<T>& operator=(const ::tools::safe_int<T>& other) {
      this->m_type = other.m_type;
      this->m_value = other.m_value;
      return *this;
    }

    static constexpr ::tools::safe_int<T> infinity() {
      return tools::safe_int<T>(::tools::safe_int<T>::type::pos_inf);
    }
    static constexpr ::tools::safe_int<T> nan() {
      return tools::safe_int<T>(::tools::safe_int<T>::type::nan);
    }

  private:
    static constexpr int f1(const ::tools::safe_int<T>& n) {
      switch (n.m_type) {
      case ::tools::safe_int<T>::type::finite:
        return 0;
      case ::tools::safe_int<T>::type::pos_inf:
        return 1;
      default: // nan
        return 2;
      }
    };
    static constexpr int f2(const ::tools::safe_int<T>& n) {
      switch (n.m_type) {
      case ::tools::safe_int<T>::type::finite:
        if (n.m_value == 0) {
          return 0;
        } else {
          return 1;
        }
      case ::tools::safe_int<T>::type::pos_inf:
        return 2;
      default: // nan
        return 3;
      }
    };
    static constexpr ::std::optional<::tools::safe_int<T>> Q() {
      return ::std::nullopt;
    }
    static constexpr ::std::optional<::tools::safe_int<T>> Z() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(0));
    }
    static constexpr ::std::optional<::tools::safe_int<T>> P() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(::tools::safe_int<T>::type::pos_inf));
    }
    static constexpr ::std::optional<::tools::safe_int<T>> U() {
      return ::std::optional<::tools::safe_int<T>>(::tools::safe_int<T>(::tools::safe_int<T>::type::nan));
    }
    static constexpr ::std::optional<bool> BQ() {
      return ::std::nullopt;
    }
    static constexpr ::std::optional<bool> BF() {
      return ::std::optional<bool>(false);
    }
    static constexpr ::std::optional<bool> BT() {
      return ::std::optional<bool>(true);
    }

  public:
    constexpr bool is_finite() const {
      return this->m_type == ::tools::safe_int<T>::type::finite;
    }

    constexpr bool is_nan() const {
      return this->m_type == ::tools::safe_int<T>::type::nan;
    }

    constexpr T val() const {
      assert(this->is_finite());
      return this->m_value;
    }

    friend constexpr bool operator==(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<bool>, 3>, 3>({{
        {BQ(), BF(), BF()},
        {BF(), BT(), BF()},
        {BF(), BF(), BF()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      return x.m_value == y.m_value;
    }
    friend constexpr bool operator!=(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      return !(x == y);
    }

    constexpr ::tools::safe_int<T> operator+() const {
      return *this;
    }
    constexpr ::tools::safe_int<T> operator-() const {
      constexpr auto table = ::std::array<::std::optional<::tools::safe_int<T>>, 3>({
        {Q(), U(), U()}
      });
      if (const auto r = table[f1(*this)]; r) return *r;

      if (this->m_value > 0) return *U();
      return ::tools::safe_int<T>(0);
    }

    friend constexpr ::tools::safe_int<T> operator+(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 3>, 3>({{
        {Q(), P(), U()},
        {P(), P(), U()},
        {U(), U(), U()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      if (y.m_value > 0 && x.m_value > ::std::numeric_limits<T>::max() - y.m_value) return *U();
      return ::tools::safe_int<T>(x.m_value + y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator+(const ::tools::safe_int<T>& x, const T& y) {
      return x + tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator+(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) + y;
    }

    friend constexpr ::tools::safe_int<T> operator-(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 3>, 3>({{
        {Q(), U(), U()},
        {P(), U(), U()},
        {U(), U(), U()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      if (x.m_value < y.m_value) return *U();
      return ::tools::safe_int<T>(x.m_value - y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator-(const ::tools::safe_int<T>& x, const T& y) {
      return x - tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator-(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) - y;
    }

    friend constexpr ::tools::safe_int<T> operator*(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 4>, 4>({{
        {Z(), Z(), U(), U()},
        {Z(), Q(), P(), U()},
        {U(), P(), P(), U()},
        {U(), U(), U(), U()}
      }});
      if (const auto r = table[f2(x)][f2(y)]; r) return *r;

      if (x.m_value > ::std::numeric_limits<T>::max() / y.m_value) {
        return *U();
      }
      return ::tools::safe_int<T>(x.m_value * y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator*(const ::tools::safe_int<T>& x, const T& y) {
      return x * tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator*(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) * y;
    }

    friend constexpr ::tools::safe_int<T> operator/(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 4>, 4>({{
        {U(), Z(), Z(), U()},
        {U(), Q(), Z(), U()},
        {U(), P(), U(), U()},
        {U(), U(), U(), U()}
      }});
      if (const auto r = table[f2(x)][f2(y)]; r) return *r;

      return ::tools::safe_int<T>(x.m_value / y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator/(const ::tools::safe_int<T>& x, const T& y) {
      return x / tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator/(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) / y;
    }

    friend constexpr ::tools::safe_int<T> operator%(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<::tools::safe_int<T>>, 4>, 4>({{
        {U(), Z(), U(), U()},
        {U(), Q(), U(), U()},
        {U(), U(), U(), U()},
        {U(), U(), U(), U()}
      }});
      if (const auto r = table[f2(x)][f2(y)]; r) return *r;

      return ::tools::safe_int<T>(x.m_value % y.m_value);
    }
    friend constexpr ::tools::safe_int<T> operator%(const ::tools::safe_int<T>& x, const T& y) {
      return x % tools::safe_int<T>(y);
    }
    friend constexpr ::tools::safe_int<T> operator%(const T& x, const ::tools::safe_int<T>& y) {
      return tools::safe_int<T>(x) % y;
    }

    constexpr ::tools::safe_int<T>& operator+=(const ::tools::safe_int<T>& other) {
      return *this = *this + other;
    }
    constexpr ::tools::safe_int<T>& operator+=(const T& other) {
      return *this = *this + ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator-=(const ::tools::safe_int<T>& other) {
      return *this = *this - other;
    }
    constexpr ::tools::safe_int<T>& operator-=(const T& other) {
      return *this = *this - ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator*=(const ::tools::safe_int<T>& other) {
      return *this = *this * other;
    }
    constexpr ::tools::safe_int<T>& operator*=(const T& other) {
      return *this = *this * ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator/=(const ::tools::safe_int<T>& other) {
      return *this = *this / other;
    }
    constexpr ::tools::safe_int<T>& operator/=(const T& other) {
      return *this = *this / ::tools::safe_int<T>(other);
    }
    constexpr ::tools::safe_int<T>& operator%=(const ::tools::safe_int<T>& other) {
      return *this = *this % other;
    }
    constexpr ::tools::safe_int<T>& operator%=(const T& other) {
      return *this = *this % ::tools::safe_int<T>(other);
    }

    constexpr ::tools::safe_int<T>& operator++() {
      return *this += ::tools::safe_int<T>(T(1));
    }
    constexpr ::tools::safe_int<T> operator++(int) {
      const auto r = *this;
      ++(*this);
      return r;
    }
    constexpr ::tools::safe_int<T>& operator--() {
      return *this -= ::tools::safe_int<T>(T(1));
    }
    constexpr ::tools::safe_int<T> operator--(int) {
      const auto r = *this;
      --(*this);
      return r;
    }

    friend constexpr bool operator<(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<bool>, 3>, 3>({{
        {BQ(), BT(), BF()},
        {BF(), BF(), BF()},
        {BF(), BF(), BF()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      return x.m_value < y.m_value;
    }
    friend constexpr bool operator>(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      constexpr auto table = ::std::array<::std::array<::std::optional<bool>, 3>, 3>({{
        {BQ(), BF(), BF()},
        {BT(), BF(), BF()},
        {BF(), BF(), BF()}
      }});
      if (const auto r = table[f1(x)][f1(y)]; r) return *r;

      return x.m_value > y.m_value;
    }
    friend constexpr bool operator<=(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      return x < y || x == y;
    }
    friend constexpr bool operator>=(const ::tools::safe_int<T>& x, const ::tools::safe_int<T>& y) {
      return x > y || x == y;
    }

    friend ::std::istream& operator>>(::std::istream& is, ::tools::safe_int<T>& self) {
      self.m_type = ::tools::safe_int<T>::type::finite;
      return is >> self.m_value;
    }
    friend ::std::ostream& operator<<(::std::ostream& os, const ::tools::safe_int<T>& self) {
      switch (self.m_type) {
      case ::tools::safe_int<T>::type::finite:
        return os << self.m_value;
      case ::tools::safe_int<T>::type::pos_inf:
        return os << "inf";
      default: // nan
        return os << "nan";
      }
    }
  };
}


#line 8 "tools/floor_kth_root.hpp"

namespace tools {

  template <typename T, typename U>
  T floor_kth_root(const T x, const U k) {
    assert(x >= 0);
    assert(k >= 1);

    if (k == 1) return x;
    if (k == 2) return ::tools::floor_sqrt(x);
    if (k == 3) {
      T ok = 0;
      T ng;
      for (ng = 1; ng * ng <= x / ng; ng *= 2);

      while (ng - ok > 1) {
        const T mid = ok + (ng - ok) / 2;
        if (mid * mid <= x / mid) {
          ok = mid;
        } else {
          ng = mid;
        }
      }
      return ok;
    }

    T ok = 0;
    T ng;
    for (ng = 1; ::tools::pow(::tools::safe_int<T>(ng), k) <= ::tools::safe_int<T>(x); ng *= 2);

    while (ng - ok > 1) {
      const T mid = ok + (ng - ok) / 2;
      if (::tools::pow(::tools::safe_int<T>(mid), k) <= ::tools::safe_int<T>(x)) {
        ok = mid;
      } else {
        ng = mid;
      }
    }

    return ok;
  }
}


#line 5 "tests/floor_kth_root.test.cpp"

using ull = unsigned long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int T;
  std::cin >> T;
  for (int t = 0; t < T; ++t) {
    ull A;
    int K;
    std::cin >> A >> K;
    std::cout << tools::floor_kth_root(A, K) << '\n';
  }

  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	all_k2_00	AC	201 ms	4 MB
g++	all_k2_01	AC	205 ms	4 MB
g++	all_k3_00	AC	122 ms	4 MB
g++	all_k3_01	AC	121 ms	4 MB
g++	all_k3_2_00	AC	166 ms	4 MB
g++	all_k3_2_01	AC	166 ms	4 MB
g++	example_00	AC	6 ms	4 MB
g++	near_border_00	AC	408 ms	4 MB
g++	near_border_01	AC	409 ms	4 MB
g++	near_border_02	AC	414 ms	4 MB
g++	near_border_2_00	AC	645 ms	4 MB
g++	near_border_2_01	AC	643 ms	4 MB
g++	near_border_2_02	AC	644 ms	4 MB
g++	random_00	AC	659 ms	4 MB
g++	random_01	AC	784 ms	4 MB

Back to top page