I want to find the greatest integer less than or equal to the kth root of n. I tried
int(n**(1/k))
But for n=125, k=3 this gives the wrong answer! I happen to know that 5 cubed is 125.
>>> int(125**(1/3))
4
What's a better algorithm?
Background: In 2011, this slip-up cost me beating Google Code Jam problem Expensive Dinner.
One solution first brackets the answer between lo and hi by repeatedly multiplying hi by 2 until n is between lo and hi, then uses binary search to compute the exact answer:
def iroot(k, n):
hi = 1
while pow(hi, k) < n:
hi *= 2
lo = hi // 2
while hi - lo > 1:
mid = (lo + hi) // 2
midToK = pow(mid, k)
if midToK < n:
lo = mid
elif n < midToK:
hi = mid
else:
return mid
if pow(hi, k) == n:
return hi
else:
return lo
A different solution uses Newton's method, which works perfectly well on integers:
def iroot(k, n):
u, s = n, n+1
while u < s:
s = u
t = (k-1) * s + n // pow(s, k-1)
u = t // k
return s
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With