Kth Number Generator in Python

July 16, 2013

A while back we looked at how to generate prime numbers in python. I was looking around for some similar problems that uses prime numbers and I came across an interview question. The question goes something lie this:

        Assume you have a function

f(x,y,z) = a^x * b^y * c^z


a,b,c - constants
x,y,z - larger than zero integers

The problem: find the kth number this function generates.

I do realize that this sounds very abstract. Let me simplify it. Assume a,b,c are 2,3,5 respectfully. If we start plugging in values to x,y,z starting from (0,0,0) we will get the following result:

 2^0 * 3^0 * 5^0 = 1
 2^1 * 3^0 * 5^0 = 2
 2^0 * 3^1 * 5^0 = 3
 2^1 * 3^1 * 5^0 = 6
 2^0 * 3^0 * 5^1 = 5
 2^1 * 3^0 * 5^1 = 10
 2^0 * 3^1 * 5^1 = 15
 2^1 * 3^1 * 5^1 = 30
 2^2 * 3^0 * 5^0 = 4

You will notice that the numbers generated are not in order. That means that brute force is almost for sure out of the question. Why? Because there is always a slight chance that we did not generate a number. For example, if we considered only the values 0,1 for x,y,z, we will never get the value 4. So what can we do? We need to think. Keep in mind that this is an interview question for programmers and mathematicians. You may see it in the future. When you do see it, you have about 30 minutes to figure it out. When you are in that scenario, do not worry about getting to the perfect solution (unless it is required). The interview is looking for a process, a thinking mind. So think out loud.

Let us look at a brute force approach to this problem:

import sys

def main(argv):

if len(argv) != 4:
sys.exit('Usage: <a> <b> <c> <k>')

a = int(sys.argv[1])
b = int(sys.argv[2])
c = int(sys.argv[3])
k = int(sys.argv[4])
list = []

for x in range(k):
for y in range(k):
for z in range(k):
list.append((a**x * b**y * c**z))

print '\nThe kth number is: ' + str(list[m-1]) + '\n'

if __name__ == "__main__":

And the output:


Now that we saw a solution, let us try to do better. After all, this solution will take a while to run for a large k. In addition, it has the chance to miss a number because we are not generating the numbers in order. If you sit for long enough and keep generating numbers on paper you may start to notice something. The numbers the function generate do have some basic factors that are easy to miss. They are products of constant (prime numbers in his case). That means that we can write the numbers we get as factors:


How about that. What if instead of trying to generate the next number based on the exponent values, we looked at what would be the minimum value? That will save us from sorting in the end and we won’t have to generate a whole bunch of extra numbers. To further use this idea, let us create 3 queues for each of the constant. Every iteration we will add or remove from the queue values based on the current value. here is a general algorithm:

        let q2,q3,q5 be queues
add 1 to q2
set val to 0
for i in k (k being the kth number)
for each q
if q_m is not empty, pop q_m to v_m
else, set v_m to MAX_INT
set val to the min value of v2,v3,v5

if val is from q2
append (2*val) to q2
append (3*val) to q3
elif val is from q3
append (3*val) to q3

append (5*val) to q5 (always)

return val

Seems simple right? Here is the code in Python:

import sys
from collections import deque

def main(argv):

if len(argv) != 5:
sys.exit('Usage: <a> <b> <c> <k> <echo>')

# get bases
a = int(sys.argv[1])
b = int(sys.argv[2])
c = int(sys.argv[3])
k = int(sys.argv[4])
echo = int(sys.argv[5])

# setup queue
q1 = deque([])
q2 = deque([])
q3 = deque([])

# init variables
val = 0

for i in range(k):

# set v to the next value in queue or to MAX_INT if queue empty
if len(q1) > 0: v1 = q1[0]
else: v1 = 2**32

if len(q2) > 0: v2 = q2[0]
else: v2 = 2**32

if len(q3) > 0: v3 = q3[0]
else: v3 = 2**32

# choose the next minimum value from the 3 queues
val = min(v1,v2,v3)

# add next values to queue
if val == v1:
elif val == v2:
elif val == v3:

# always add the largest

# if echo is True print every number
if echo: print str(i+1) + ": " + str(val)

print '\nThe kth number is: ' + str(val) + '\n'

if __name__ == "__main__":

And here are some example runs:



If you enjoyed this post, please consider leaving a comment or subscribing to the RSS feed to have future articles delivered to your feed reader.

You Might Also Like