From afe1859bc6a5de62907987348263d0a2a7238a8c Mon Sep 17 00:00:00 2001 From: Matej Focko Date: Sat, 12 Aug 2023 23:20:17 +0200 Subject: [PATCH] chore(cpp): implement isqrt Signed-off-by: Matej Focko --- .common/cpp/skeleton.cpp | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/.common/cpp/skeleton.cpp b/.common/cpp/skeleton.cpp index a8686dd..01588ab 100644 --- a/.common/cpp/skeleton.cpp +++ b/.common/cpp/skeleton.cpp @@ -158,6 +158,28 @@ long pow(long base, long exp) { return half * half * base; } +template +T isqrt(T x) { + assert(x >= 0); + + auto max_shift = 8 * sizeof(T) - 1; + auto shift = (max_shift - std::countl_zero(x)) & ~1; + auto bit = 1 << shift; + + T result = 0; + while (bit != 0) { + if (x >= (result + bit)) { + x -= result + bit; + result = (result >> 1) + bit; + } else { + result = (result >> 1); + } + bit = bit >> 2; + } + + return result; +} + template struct Z { Z(std::int64_t x = 0) : x(x % MODULO) {}