pythonalgorithmmathnth-root

How to find integer nth roots?


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.


Solution

  • 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