reservoir sampling of k items

Given a stream of n elements where n may be very large or unknown, randomly select exactly k elements such that every element has equal probability of being chosen.

The solution must process the stream in a single pass using only O(k) extra space.

import random
def reservoir_sampling_k(stream, k):
    reservoir = []
    i = 0

    for val in stream:
        i += 1

        if i <= k: #first K numbers should get the chance to be chosen first
            reservoir.append(val)
        else: #now if current index is more than k then take a random guess from 1 to i, if the random guess is less than equal to k then push the current element in that index
            j = random.randint(1, i)

            if j <= k:
                reservoir[j - 1] = val  # replace

    return reservoir