Choose K of the First N Numbers
Problem
Given two integers n and k, return every way to choose k different numbers from 1 to n. Order inside a choice does not matter, so write each choice in ascending order, and return the choices in ascending lexicographic order.
Examples
Input: n = 4, k = 2
Output: [[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
Why: 4 choose 2 is 6Input: n = 3, k = 3
Output: [[1, 2, 3]]
Why: there is only one way to take everythingInput: n = 2, k = 0
Output: [[]]
Why: edge case, choosing nothing is exactly one choice, the empty oneHints
0 / 3
Build a choice one number at a time. To keep each choice in ascending order and avoid producing [2, 1] after [1, 2], only ever add numbers larger than the last one you added.
Recurse with a start value: pick a number x from start upward, recurse with start x + 1, then remove x and try the next one. Record a copy when the choice has k numbers.
Prune the loop: if you still need m more numbers, starting at x only works when x + m - 1 <= n, so the loop can stop at n - m + 1 instead of n.
Solution
This is the subsets decision tree cut off at depth k. Each level decides the next number of the choice, and passing x + 1 as the next start guarantees every choice is built in ascending order, so each set of numbers is produced exactly once and the output comes out in lexicographic order for free. After the recursive call returns, the number is popped, which puts the shared pick list back the way the loop found it. The upper bound on the loop is the pruning step: a branch that cannot possibly collect enough numbers is never entered. There are C(n, k) choices and each costs O(k) to copy, so time is O(k × C(n, k)) and the recursion depth is k.
def choose(n, k):
out, pick = [], []
def build(start):
if len(pick) == k:
out.append(pick[:]) # record a copy, pick keeps changing
return
need = k - len(pick)
for x in range(start, n - need + 2): # leave room for the numbers still needed
pick.append(x)
build(x + 1)
pick.pop() # undo before trying the next x
build(1)
return out
print(choose(4, 2)) # -> [[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
print(choose(3, 3)) # -> [[1, 2, 3]]
print(choose(2, 0)) # -> [[]]Stuck on the idea rather than the code? Subsets covers it.