Check if a number N is a power of K

Question

I was asked this question in interview:

Check if a number N is a power of K.

Example: N = 32, K = 2 => True

N = 40, K = 5 => False

I wrote following code but got the feedback that, complexity could have be improved, How to improve its complexity?

def check_kth_power(n, k):
        while n%k == 0:
            n = n/k

        if n != 1:
            return False
        return True

        
print check_kth_power(128, 5)

Answer 1

You could have used the fact that if \$n = k^x\$, then \$\log_k(n) = x\$ is an integer. However, due to the limited precision of floats in Python, this will generate some problems. To get around this, we can test both the ceil and the floor of the logarithm to see if we can get back n. This way we only ever need to do three expensive computations at most.

from math import log, ceil, floor

def check_kth_power(n, k):
    kth_log = log(n, k)
    return k ** int(ceil(kth_log)) == n or \
           k ** int(floor(kth_log)) == n

Now, instead of doing the exponentiation twice, we can realize that k ** int(ceil(kth_log)) = k * k ** int(floor(kth_log)) for a small speed-boost in about 50% of the cases (thanks to @KyleGullion:

def check_kth_power(n, k):
    candidate = k ** int(log(n, k))
    return candidate == n or k * candidate == n

This seems to work for quite large values (I checked up to about \$2^{32000}, 2^{20000} - 1, 3^{20000}\$). This is the version in the timing table below.

To get even more of a speed-boost, you can steal the nice premature exit from @pantarei's answer. This sacrifices some speed in the case that the number is a power of k, but gains about a factor 2 if it is not:

def check_kth_power_gp(n, k):
    if n%k != 0:
        return False
    candidate = k ** int(log(n, k))
    return candidate == n or k * candidate == n

---

And, finally, some timings. This compares my code and the code from @user1825464's answer, @pantarei's answer, @RyanMill's answer, @Eevee's answer and @JamesMishra's answer:

+-------------+---+----------+-------------+-----------+------------+---------+--------------+
|      n      | k | Graipher | user1825464 | panta rei | Ryan Mills |  Eevee  | James Mishra |
+-------------+---+----------+-------------+-----------+------------+---------+--------------+
| 2**1000     | 2 | 1.73 µs  | 9.9 µs      | 105 µs    | 12.9 µs    | 388 µs  | 3.23 µs      |
| 2**1000 - 1 | 2 | 1.99 µs  | 2.5 µs      | 619 ns    | 15.7 µs    | 765 ns  | 3.09 µs      |
| 3**1000     | 2 | 2.41 µs  | 4.26 µs     | 854 ns    | 22.4 µs    | 1.04 µs | 4.08 µs      |
| 3**1000     | 3 | 2.81 µs  | 12.6 µs     | 125 µs    | 13.8 µs    | 556 µs  | 4.51 µs      |
+-------------+---+----------+-------------+-----------+------------+---------+--------------+

So the log does not care if it actually is a kth power, whereas the loop has to do more work in that case, but potentially finds that it is not a power faster.

---

You could be checking for integerness using x % 1 == 0 for integers, != 0 otherwise (thanks to @Peilonrayz in the comments):

from math import log

def check_kth_power(n, k):
    return log(n, k) % 1 == 0

But, as noted in the comments, this works only until the precision of float is not enough anymore to distinguish that the result of the log is not an integer.

For numbers of the form \$ 2^x - 1\$, this happens at \$x = 48\$, as noted by @GarethReese, so you should have asked if there are any limits on n and k.

Answer 2

def check_kth_power(n, k):
    if n == 1:
        return True

    div = k
    while div * div <= n:
        div = div * div

    if n % div != 0:
        return False
    return check_kth_power(n / div, k)

The complexity of this is \$O(log(log(N))\$, compared to \$ O(log(N)) \$ for straightforward division or logarithms.

The reasoning behind this is that if \$ N \$ is a power of \$ K \$, its representation in base \$K\$ looks something like like \$000100 ..00\$. We can do something like a binary search to find the position of this one (if it exists) by checking the largest divisor of \$N\$ of the form \$ K^{2^i} \$, dividing it out, then doing this recursively until it either doesn't divide out (and isn't a power of \$K\$), or it returns 1 (and is a power of \$K\$).

Answer 3

Bit of a wild guess, but perhaps they were looking for you to use divmod, since % and / with the same operands are essentially the same operation?

def check_kth_power(n, k):
    while True:
        q, r = divmod(n, k)
        if r == 0:
            n = q
        else:
            break

    return n == 1

Answer 4

This is a version of @Graipher's answer, but it relies more on properties of Python ints.

from math import log
def check_kth_power(n, k):
    test_pow = int(log(n, k))
    test_n_floor = k ** test_pow
    test_n_ceil = k ** (test_pow + 1)
    return n == test_n_floor or n == test_n_ceil

If test_pow suffers from floating point issues, then neither test_n_floor nor test_n_ceil will equal n. This is because the Python exponent operator on two ints will produce either an int or long type, and will not lose precision.

>>> log(2 ** 50, 2)
50.0
>>> log(2 ** 50 - 1, 2)
50.0
>>> check_kth_power(2 ** 50, 2)
True
>>> check_kth_power(2 ** 50 - 1, 2)
False
>>> 2 ** 50 - 1 == 2 ** int(log(2 ** 50-1, 2))
False
>>> check_kth_power(3 ** 5, 3)
True
>>> check_kth_power(3 ** 5 + 1, 3)
False
>>> check_kth_power(3 ** 5 - 1, 3)

False

Credit goes to @Graphier for discovering a bug in my original answer. I have updated with a fix.

Answer 5

Why not doing it the (straight) forward way?

def check_kth_power(n, k):
    if n%k != 0:
        return False

    pow = k
    while pow < n:
        pow = pow * k
    return pow == n

print check_kth_power(128, 5)

This code is easier to understand, and uses more efficient operations.

But @Graipher's answer is probably the most efficient way to do it.

Answer 6

Binary search

They're interested in complexity. I'm pretty sure they wanted you to use binary search.

This is an interview question. So you have to show what you know about algorithms and complexity.

• Just doing exponentiation by repeated multiplication (or division) doesn't cut it.

• Just using math.log, which is probably how you would do it in practice, also won't get your knowledge across.

To do binary search, first you need a fast way to get an upper bound. Then implement binary search in the standard textbook way, by shifting the upper (hi) and lower (lo) bounds.

def got_the_power(n, k):
    '''
    Returns True if n is a power of k, otherwise returns False.    
    '''

    def get_upper_bound(n, k):
        '''
        Finds an upper bound for e where k**e = n.
        '''
        a, e = k**2, 2
        while a < n:
            a *= a
            e = e*2
              # The above line was e = e**2 in the original post.
              # Bug spotted by Graipher. Thanks! 
        return e

    hi = get_upper_bound(n, k)
    lo = 0
    while hi - lo > 1:    # standard binary search
        m = (hi - lo) / 2 + lo
        a = k**m
        if a == n:
            return True
        if a > n:
            hi = m
        else:
            lo = m
    return False

# Test it out
print got_the_power(64, 2)  # True
print got_the_power(63, 2)  # False

Both the upper bound subroutine and the overall function each run in \$\mathcal{O}(log(e))\$., where e is the log of n to base k. However, the fundamental operation is exponentiation, which is itself in \$\mathcal{O}(log(e))\$. So the overall complexity is \$\mathcal{O}(log(e)^2)\$.

Answer 7

Complexity of your code is very, very bad if n == 0 or k == 1 because your loop doesn't finish :-( So that should be fixed.

When people discuss complexity, there's worst case and average case. In computer science, people often look primarily at the worst case. In real life the average case is more important.

In your case, the worst case is that n = k^m, in which case it takes m steps. However, if n is chosen at random, then the chances are (k-1) / k that you stop after just one step, (k-1) / k^2 for two steps, (k-1) / k^3 for three steps and so on. Even for k = 2, the average number of steps is just two.

Of course that's just me saying that your code is Ok. In a job interview, you may have to live with whatever the interviewer says.