Itertools Recipe: Power Set

The itertools docs has a ton of slick recipes for using the library to good effect. Some of the code is more useful than illustrative, so I wanted to use these notebooks to break down a few of the functions.

This is

# poor import style, but I want to copy-paste the code
# as-is from the docs

from itertools import *
import itertools

power_set()

This one’s one of my favorite recipes.

The Power Set is basically a set that exhaustively says “what would this group look like if this element were/weren’t in the group”… but for every possible combination. The resulting set winds up containing 2^n elements, as each element can “either be in the set, or not”.

As you might imagine, this could lead to a ton of messy, complicated flow-control.

Or a slick one-liner using itertools

def power_set(iterable):
    "powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3)"
    s = list(iterable)
    return chain.from_iterable(combinations(s, r) for r in range(len(s)+1))

Demo

letters = ['a', 'b', 'c', 'd']
for group in power_set(letters):
    print(group, end=' ')
() ('a',) ('b',) ('c',) ('d',) ('a', 'b') ('a', 'c') ('a', 'd') ('b', 'c') ('b', 'd') ('c', 'd') ('a', 'b', 'c') ('a', 'b', 'd') ('a', 'c', 'd') ('b', 'c', 'd') ('a', 'b', 'c', 'd') 

Why this works

It helps to read this one inside-out, I think.

The Generator Comprehension

First, notice that the generator comprehension is executing over ... for in range(len(s) + 1). This step is feeding the r that’s used in combinations(). The second argument is “return all possible combinations of the original set s of r elements”.

So we obviously want to start with 0 (no elements are in the set). Then move on to what the set looks like when only one of each element are in it. Then all combinations of 2, etc, etc.

Thus, the last set generated should be len(s)==r, where every element is returned as in the set.

However, because the range operator will stop at n-1 (e.g. range(4) prints 0 1 2 3), we add 1 to the length of our list s so that it’s the same.

Chaining

Running the function by itself yields a chain object

power_set(letters)
<itertools.chain at 0x277f6030be0>

This is because we’re still lazily yielding the values generated by the function. chain.from_iterable() is going to assemble one big iterable from arbitrarily-many little ones.

Can’t spitball an immediate use for this just yet, but this basically means that we can define the Power Set once, and pass it around, iterating to get the next combinations as we need them.

ourset = power_set(letters)
for i in range(4):
    print(next(ourset), end=' ')
() ('a',) ('b',) ('c',) 
for i in range(4):
    print(next(ourset), end=' ')
('d',) ('a', 'b') ('a', 'c') ('a', 'd') 
for i in range(4):
    print(next(ourset), end=' ')
('b', 'c') ('b', 'd') ('c', 'd') ('a', 'b', 'c') 
for i in range(4):
    print(next(ourset), end=' ')
('a', 'b', 'd') ('a', 'c', 'd') ('b', 'c', 'd') ('a', 'b', 'c', 'd')