Check if a number N is a power of K

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)

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.

StackExchange Code Review Q#160924, answer score: 17