Compute a random subset - Python

random_subset.py

# This solution to compute a random subset avoids the use of extra auxiliary space, with O(k + n) time complexity and O(1) space complexity, in Python (5.15 Compute a random subset, on EPI (Elements of Programming Interviews)) (September 2018 edition)).
# The idea here is to pick the r-th combination, and find its constituents by incrementing them in a Odometer like fashion, and taking into account that the next digit in the combination will be greater than the previous one.
# For example, the combination sequence for n=5 and k=2 is:
# 0 - [0,1]
# 1 - [0,2]
# 2 - [0,3]
# 3 - [0,4]
# 4 - [1,2]
# 5 - [1,3]
# 6 - [1,4]
# 7 - [2,3]
# 8 - [2,4]
# 9 - [3,4]
# To know how many combinations exist for a given n/k or how many combinations exist until the next digit change, the following formula is used: [n!/((n - k)!k!)](https://en.wikipedia.org/wiki/Combination)


import random, math

def random_subset(n, k):
    def num_combinations(n0, k0):
        return math.factorial(n0) // (math.factorial(n0 - k0) * math.factorial(k0))

    combination_idx = random.randint(0, num_combinations(n, k) - 1)
    result = [0] * k
    counter = 0

    for k_i in range(k):
        if (k_i > 0):
            result[k_i] = result[k_i - 1] + 1

        while (n - 1 - result[k_i] >= 0 
               and combination_idx >= counter + num_combinations(n - 1 - result[k_i], k - k_i - 1)):
            counter += num_combinations(n - 1 - result[k_i], k - k_i - 1)
            result[k_i] += 1

    return result
  
# Example
for i in range(100):
    print(random_subset(5, 3))