Subsets with Bitmasks

20 min · Intermediate

A set of n items has 2 to the power n subsets. A bitmask gives each subset a number: bit i says whether item i is in. Counting from 0 up to 2**n - 1 then visits every subset once, with no recursion and no extra lists.

A mask is a yes or no list

A bitmask is an integer used as a row of yes or no flags, one per bit. For n items, each number from 0 to 2**n - 1 stands for exactly one subset.

Take three items. A 3-bit number says, for each item, in or out. The number 0b101 means "item 0 and item 2 are in, item 1 is out".

Python

items = ["a", "b", "c"]mask = 0b101print([items[i] for i in range(3) if (mask >> i) & 1])

The comprehension keeps item i when bit i of the mask is 1. The result is ['a', 'c'].

Visit every subset

There are 1 << n masks from 0 to 2**n - 1, so one loop over range(1 << n) covers every subset, including the empty one (mask 0) and the full one.

Python

items = ["a", "b", "c"]n = len(items)for mask in range(1 << n):    print(mask, [items[i] for i in range(n) if (mask >> i) & 1])

Example 1 · The same subsets with itertools

Python

from itertools import combinations
items = [1, 2, 3]subsets = [list(c) for size in range(len(items) + 1) for c in combinations(items, size)]print(subsets)

Output

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

What this code does

combinations yields the subsets of each size. Joined for sizes 0 to n it gives all 2**n subsets.
The order differs from the mask loop (by size, not by number) but the subsets are the same.

Example 2 · How fast the count grows

Python

print(1 << 10, 1 << 20, 1 << 30)

Output

1024 1048576 1073741824

What this code does

Each extra item doubles the number of subsets.
About a million subsets at n = 20 is fine in Python. A billion at n = 30 is not, so masks are for small n.

It prints 8 lines. Mask 0 gives [] and mask 7 gives all three items. The cost is O(2*n n), fine up to about n = 20.

Subset sum: try every mask

A shop has gift cards worth 3, 5 and 8, and wants to know in how many ways they add up to exactly 8. Try every subset and count the matches.

Python

1prices = [3, 5, 8]2target = 83n = len(prices)4count = 05for mask in range(1 << n):6    total = sum(prices[i] for i in range(n) if (mask >> i) & 1)7    if total == target:8        count += 19print(count)

Example 3 · Split a team evenly

Python

1weights = [3, 1, 4, 2, 2]2total = sum(weights)3best = total4for mask in range(1 << len(weights)):5    s = sum(w for i, w in enumerate(weights) if (mask >> i) & 1)6    best = min(best, abs(total - 2 * s))7print(best)

Output

What this code does

total - 2 * s is the gap between the two groups when one group sums to s.
The smallest gap is 0: 4 + 2 on one side and 3 + 1 + 2 on the other.

Example 4 · Subsets grouped by size

Python

from collections import Counter
sizes = Counter(mask.bit_count() for mask in range(8))print(sorted(sizes.items()))

Output

[(0, 1), (1, 3), (2, 3), (3, 1)]

What this code does

The number of set bits of a mask is the size of its subset.
Three items give 1 empty set, 3 single items, 3 pairs and 1 full set.

The subsets {3, 5} and {8} both reach 8, so it prints 2. For large n a smarter method is needed, but this brute force is the base to check it against.

Submasks and visited sets

Sometimes you need only the subsets inside one mask. The step sub = (sub - 1) & mask jumps to the next smaller submask. A whole visited set can also live in one integer.

Python

1mask = 0b10112sub = mask3found = []4while sub:5    found.append(bin(sub))6    sub = (sub - 1) & mask7print(found)

Example 5 · A visited set in one integer

Python

1visited = 02for city in (2, 0, 2, 5):3    if visited & (1 << city):4        print("already", city)5    visited |= 1 << city6print(bin(visited), visited.bit_count())

Output

already 2
0b100101 3

What this code does

Bit i of visited records city i. AND tests it, OR sets it.
City 2 comes twice, so the second visit is reported. Three different cities were seen.

It lists the 7 non-empty submasks of 1011, from largest to smallest. Over all masks of n bits the total work is 3n, not 4n.

Next step

Check your understanding

10 quick questions

1prices = [3, 5, 8]2target = 83n = len(prices)4count = 05for mask in range(1 << n):6 total = sum(prices[i] for i in range(n) if (mask >> i) & 1)7 if total == target:8 count += 19print(count)

1weights = [3, 1, 4, 2, 2]2total = sum(weights)3best = total4for mask in range(1 << len(weights)):5 s = sum(w for i, w in enumerate(weights) if (mask >> i) & 1)6 best = min(best, abs(total - 2 * s))7print(best)