kth-prime.py

import math
def kthPrime(k):
    # k is non-negative
    if k < 0:
        return None
    # an array to store primes
    primes = []
    # use an array of booleans `isPrime` to indicate whether a number is a prime
    # e.g. isPrime[2] is True means 2 is a prime
    # upperLimit is the size of the array. We start from a small number
    # and it becomes bigger when necessary
    upperLimit = 16
    isPrime = [True] * upperLimit
    isPrime[0], isPrime[1] = False, False # 0 and 1 are not used
    # prime numbers start from 2
    start = 2
    while len(primes) <= k:
        # double the size of isPrime
        isPrime += [True] * upperLimit
        upperLimit *= 2
        for i in xrange(2, int(math.sqrt(upperLimit)) + 1):
            if isPrime[i]:
                for j in xrange(i * i, upperLimit, i):
                    isPrime[j] = False
        for i in xrange(start, upperLimit):
            if isPrime[i]:
                primes.append(i)
                if len(primes) > k:
                    break
        if len(primes) <= k:
            # next loop start from the current upperLimit
            start = upperLimit

    return primes[k]

for i in range(50):
    print kthPrime(i)
print kthPrime(499)