K-subsets in Python

Why this matters

Many everyday questions ask for a group of a fixed size. Which three people can form a committee from a team of five? Which two toppings can go on a pizza from a list of six? Which two products in a shop cost less than a given amount together? Each answer is a k-subset: a group of exactly k items taken from a list of n items. The order inside the group does not matter. Picking Ana and Ben is the same group as picking Ben and Ana.

Programmers also call this a combination. A combination is the same idea: a way to pick k things from n without caring about the order. The name you use matters less than the method. The method is backtracking, which you met in the subsets and permutations lesson. This lesson uses one new control, a start index, which lets you pick each group exactly once.

Words used in this lesson

A recursion is a function that calls itself on a smaller problem. A base case is the situation where the function stops calling itself. A backtracking function builds an answer one choice at a time and takes each choice back afterwards. The start index is a number that says where the next choice may begin in the list. A duplicate is the same group appearing more than once in the output.

The list you pick from is called the input. Its items are numbered from 0 at the first item. In Python, a list such as items = [1, 2, 3, 4] has items at positions 0, 1, 2 and 3. The method append adds an item to the end of a list, and pop() removes the last item and gives it back.

The idea from a small example

Suppose you have four cards numbered 1, 2, 3 and 4, and you want to pick two of them. Try it by hand. You could write 1 and 2, then 1 and 3, then 1 and 4. Then you move on to 2 and 3, then 2 and 4, and last 3 and 4. That gives six groups.

Look at how you did it. You never wrote 2 and 1, because that is the same group as 1 and 2. You always moved forward. After you picked 1, your second card came from the cards after 1. After you picked 2, your second card came from the cards after 2. The rule is simple: each new card must come later in the list than the card before it.

That rule is the start index. It stops the same group from being built twice, once in each order.

The number of groups is written as C(n, k), read as "n choose k", and it can be worked out by hand. For four cards and two picks, the count is 4 times 3 divided by 2 times 1, which is 12 divided by 2, which is 6. That matches the six groups you wrote down.

Choose, explore, undo, with a start index

The backtracking pattern from the subsets lesson still applies. For each step you do three things in order:

1. Choose one item and add it to the group you are building.
2. Explore by calling the function again, to pick the rest of the group.
3. Undo by removing that item, so the next choice starts from the same state.

The new part is the loop. At each level the function tries every item from the start index to the end of the list. The start index for the next level is one more than the position of the item just chosen. That is what makes the rule "later items only" work in code.

The base case is reached when the group has k items. At that moment the current group is one finished answer, and you save a copy of it.

Example 1: pairs from four numbers

This first example picks groups of two numbers from [1, 2, 3, 4]. The function takes five inputs: the list, the size of the group k, the start index, and two containers. One container holds the group being built, called current. The other holds every finished group, called result.

Press Run. The output is:

[[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
count = 6

Three details in the code matter.

The line result.append(current[:]) saves a copy of the group. The slice [:] makes a new list with the same items at that moment. Without the copy, every saved group would point at the same current list, and all of them would end up empty. The Common mistakes section shows that exact result.

The line current.pop() removes the last item in the list. It is the item just added, so the list goes back to the state it had before the choice.

The loop runs from start and not from 0. That one change is what stops the same group from appearing in more than one order.

Example 2: a committee from five people

The second example uses names instead of numbers, and picks three from five. The code is the same shape. Only the types change: the items are strings and the groups are lists of strings. This shows that the function does not depend on the items being numbers.

The output is:

['Ana', 'Ben', 'Chloe']
['Ana', 'Ben', 'Dev']
['Ana', 'Ben', 'Eli']
['Ana', 'Chloe', 'Dev']
['Ana', 'Chloe', 'Eli']
['Ana', 'Dev', 'Eli']
['Ben', 'Chloe', 'Dev']
['Ben', 'Chloe', 'Eli']
['Ben', 'Dev', 'Eli']
['Chloe', 'Dev', 'Eli']
count = 10

Check the count by hand. Five people, three per committee: 5 times 4 times 3 divided by 3 times 2 times 1 is 60 divided by 6, which is 10. The printed count is 10, so the function found every committee once.

Look at the first line of output. It starts with Ana, Ben, Chloe. Every later committee that contains Ana and Ben comes after it in the list, and the committees are printed in the order the loop reaches them. The group Ben, Ana, Chloe is never printed, because Ana comes before Ben in the input, and the function only moves forward from Ana.

Example 3: repeated values in the input

The third example is harder. What if the input contains the same value twice? Suppose the input is [2, 1, 2, 3], which has two copies of the number 2. The question is which groups should appear, and how many times.

The two copies of 2 are different items in the list. They sit at different positions. So a group that uses one copy of 2 and a group that uses the other copy look the same on the screen, and the function will list both. First, sort the input so that equal values sit next to each other. The built-in sorted function returns a new sorted list. The sorted list is [1, 2, 2, 3].

Now the duplicate problem is easy to see. Whenever the loop reaches the second 2 at the same level as the first 2, it would produce the same group again. The fix is one rule: at each level of the loop, skip any item that has the same value as the item just before it, unless it is the first item the loop starts with.

The output is:

[[1, 2], [1, 3], [2, 2], [2, 3]]
count = 4

The pair [2, 2] is correct here. The input really does contain two 2s, so a group of two 2s is a valid choice. The rule removes only the repeats that come from picking the same value at the same level.

To see what the rule saves, run the same code without the if that skips items. You get this output:

[[1, 2], [1, 2], [1, 3], [2, 2], [2, 3], [2, 3]]
count = 6

The pair [1, 2] appears twice, and [2, 3] appears twice. Those are duplicates. The skip rule removes exactly those.

Why the rule uses i > start

The check is i > start, and not i > 0. A deeper loop begins at start, and the first item it tries must always be used. Take the call that has already picked 2 from position 1. Its loop starts at position 2, which also holds 2. The check i > start is false there, so the 2 is tried and [2, 2] is saved. If the check were i > 0, that 2 would be skipped because it equals the item before it, and [2, 2] would be lost.

Tracing the calls by hand

Reading code is not enough. To be sure you understand a recursive function, trace it on paper with a small input. Take [1, 2, 3] and k equal to 2. The table below shows each call, what it chooses, and what it saves. Each indented line is a deeper call.

choose(start=0), current=[]
  choose one: pick 1, call choose(start=1), current=[1]
    choose one: pick 2, call choose(start=2), current=[1, 2]  size is 2, saved [1, 2]
    choose one: pick 3, call choose(start=3), current=[1, 3]  size is 2, saved [1, 3]
    loop ends, undo 1, current=[]
  choose one: pick 2, call choose(start=2), current=[2]
    choose one: pick 3, call choose(start=3), current=[2, 3]  size is 2, saved [2, 3]
    loop ends, undo 3, current=[2]
  choose one: pick 3, call choose(start=3), current=[3]
    loop does not run, because start is 3 and there are only 3 positions (0, 1, 2)
    returns without saving

The trace shows the three saved groups, [1, 2], [1, 3] and [2, 3], in the order they are saved. It also shows why the pick of 3 at the top level saves nothing. The call with current=[3] has one item, which is less than k, and no items come after position 2.

Doing this trace once for a small input is the best way to find a bug in your own version of the function.

Checking your output

Before you trust a new version of this function, check its output against four facts.

1. The number of groups equals C(n, k). For [1, 2, 3, 4] with k equal to 2, that is 6. For [1, 2, 3, 4, 5] with k equal to 3, it is 10.
2. Every group has exactly k items.
3. Every group appears once. No two lines of output are the same group, in any order.
4. Each group keeps the input order. A group written as [2, 1] is a sign that the start index is wrong.

The count formula is simple arithmetic. For ten items with k equal to 3, you multiply 10 times 9 times 8 and divide by 3 times 2 times 1. That is 720 divided by 6, which is 120. A function that prints 120 groups for that input passes this check. The math.comb function in the standard library returns C(n, k) directly, so you can compare your count with math.comb(10, 3), which is 120.

Cost in plain numbers

Each finished group costs one copy of a list of k items. The number of groups is C(n, k), so the output grows fast as n and k grow. For n equal to 5 and k equal to 2 there are 10 groups. For n equal to 10 and k equal to 3 there are 120. For n equal to 10 and k equal to 5 there are 252. The search also makes calls that do not save anything, such as the call with current=[3] in the trace above. Those calls are the cost of the loop, and they add up as the list gets longer.

The itertools.combinations function gives the same groups as tuples, in the same order, with no code to write. It is the right choice in real programs. The function in this lesson is still worth writing once, because it is the template for problems that the library cannot solve directly, such as a group with a sum rule.

Common mistakes

Each mistake below gives a wrong answer or an error, and each one is easy to make. The outputs shown are the real outputs of these small programs.

Mistake 1: saving the group without a copy

If you write result.append(current) instead of result.append(current[:]), every saved entry points at the same list. When the search ends, that list is empty, so every entry prints as []. The program runs and prints six entries, all empty:

[[], [], [], [], [], []]

The fix is to save a copy. The slice current[:] is a new list with the same items at that moment.

Mistake 2: forgetting the undo step

Suppose the line current.pop() is deleted. The first group saved is [1, 2]. After that, the group keeps growing with items that were never removed, so the length never equals k again. The program prints only one group:

[[1, 2]]

The undo step is what lets the next branch start from the same state as the one before. Without it, the state is polluted.

Mistake 3: removing by value or by position

Python has two ways to take an item out of a list. list.remove(value) removes the first item equal to a value. list.pop(index) removes the item at a position. They are easy to mix up, because both names sound like "take it out." If you write current.pop(items[i]) in place of the undo line, the value is used as a position. For a list with two items, position 2 does not exist, and Python stops with this error:

IndexError: pop index out of range

The fix is to use current.pop() with no argument, which removes the last item. The reverse mistake is current.remove(items[i]). That does not raise an error here, because the value is in the list, but it removes the first matching value, which can be a different copy from the one just added. It also takes time to search the list. If the value is not in the list at all, Python stops with this error:

ValueError: list.remove(x): x not in list

Mistake 4: starting the next level at start or at 0

If the recursive call passes start instead of i + 1, the same item can be chosen again at the next level. The result then contains groups such as [1, 1], which is no longer a k-subset of distinct positions. If the recursive call passes 0, the next pick can be an earlier item, so the same two items show up in both orders. In both cases, the start index is the piece that must move forward.

Here is the first case, with start passed in place of i + 1, on [1, 2, 3] with k equal to 2:

The output is:

[[1, 1], [1, 2], [1, 3], [2, 2], [2, 3], [3, 3]]

The output has 6 groups, and three of them, [1, 1], [2, 2] and [3, 3], use the same item twice. A group of two distinct items should give only the three pairs from the first example.

Mistake 5: forgetting to sort before skipping repeats

The skip rule in example 3 only works when equal values sit next to each other. If sorted is missing, the input [2, 1, 2, 3] has its two 2s apart, and the check items[i] == items[i - 1] misses some duplicates. Sort first, then apply the rule.

Practice

Exercise 1 (easy). Pick groups of two from [5, 6, 7] by hand. Write each group in order, and count them. Then check the count with the formula 3 times 2 divided by 2 times 1.

Show the answer

The groups are [5, 6], [5, 7] and [6, 7], so the count is 3. The formula gives 3 times 2 divided by 2 times 1, which is 6 divided by 2, which is 3. The two counts match. math.comb(3, 2) also gives 3.

Exercise 2 (easy). How many groups of three can you make from six items? Work it out with the formula, then say how many groups the function prints for [1, 2, 3, 4, 5, 6] with k equal to 3.

Show the answer

The formula is 6 times 5 times 4 divided by 3 times 2 times 1. That is 120 divided by 6, which is 20. The function prints 20 groups, one for each. math.comb(6, 3) gives the same 20.

Exercise 3 (medium). In Example 1, change the recursive call from choose(items, k, i + 1, current, result) to choose(items, k, i, current, result). Predict the first group printed, and say whether the output still contains only groups of distinct positions. Then run it to check.

Show the answer

The first group printed is [1, 1], because the loop at position 0 picks 1 again at the next level. The output contains groups that use the same position twice, such as [1, 1] and [2, 2]. The function no longer picks k distinct items. The fix is to change i back to i + 1.

Exercise 4 (medium). Run the duplicate-skipping version from Example 3 on the input [3, 3, 3] with k equal to 2. Predict the output before you run it.

Show the answer

The output is [[3, 3]] and the count is 1. After sorting, the three 3s sit next to each other. The skip rule lets the first 3 start a group, and the second 3 is the only one allowed to fill the second position, because the third 3 is skipped there. Every other combination repeats [3, 3], so the rule removes it.

Exercise 5 (harder). Explain in your own words why the function does not need to check whether a group has already been saved. Use the example [1, 2, 3, 4] in your answer.

Show the answer

The function never builds the same group twice in the first place. Each call only uses items at later positions than the last item chosen. So the group [2, 1] cannot occur, because after choosing 2 the function only looks at 3 and 4. Each group has exactly one order it can be built in, the order of the positions, so it is saved once. The only time a repeat can appear is when the input has equal values, which is why Example 3 needs its extra rule.

Summary

• A k-subset is a group of k items from a list of n, where the order does not matter.
• A start index that moves forward after each pick stops the same group from appearing in two orders.
• The function is choose, explore, undo, with a loop from start to the end of the list.
• Save a copy of the group when it is complete, with current[:], and undo the last choice with pop() after each call returns.
• When the input has repeated values, sort it first and skip an item that equals the item just tried at the same level.
• The count is C(n, k), which you can check by hand with n times n minus 1 and so on, divided by k times k minus 1 and so on, or with math.comb.

The pruning idea, which stops a branch early when there are not enough items left to finish the group, is covered in the pruning lesson of this module. The balanced parentheses lesson uses the same backtracking pattern with a different rule for what a valid choice is.