Bitmask as a Set: Enumerate Subsets and Write Bitmask DP

When a problem has at most twenty or so items and asks about subsets of them, which ones are picked, which cities are visited, which tasks are done, there is a representation that makes the code shorter and faster at once: one integer, one bit per item. Subsets become numbers, set operations become single instructions, and "every subset" becomes a for loop. It is also the state behind a whole family of dynamic programming solutions.

The problem it solves

Many problems need to enumerate or index subsets:

• Try every combination of items and keep the best, when n is small.
• Remember which elements have been used so far, as the key of a memo table.
• Assign tasks to workers, pick a team covering all skills, or visit every city once.

A Python set or a list of booleans works, but a set cannot be a dictionary key without freezing it, and comparing or hashing a set costs time proportional to its size. An integer mask is compact, hashable, compares in one step for any realistic n, and can index a plain array directly.

The intuition

Give item i the bit 1 << i. A mask is the sum of the bits of the items it contains: with a, b, c as bits 1, 2 and 4, the mask 5 means a and c.

The set operations are the bitwise operators:

• Membership: mask >> i & 1.
• Add, remove: mask | 1 << i, mask & ~(1 << i).
• Union, intersection, difference: x | y, x & y, x & ~y. Symmetric difference is x ^ y.
• Size: the number of set bits, bin(mask).count("1") or int.bit_count() in Python 3.10 and later.

Because every subset of n items is a distinct integer from 0 to 2^n - 1, counting through that range visits each subset exactly once. No recursion, no stack of choices to undo. Two more patterns cover most interview uses:

• Submasks. sub = (sub - 1) & mask steps through every subset of a given mask, from the mask itself down to its smallest non-empty submask, and 0 is added last. Walking all submasks of all masks costs 3^n in total, not 4^n, because each item is either outside the mask, in the mask but not the submask, or in both.
• Bitmask DP. When the answer depends on which items have been used but not on the order they were used in, the mask is the DP state. The classic example is the shortest tour through all cities: best[mask][j] is the cheapest path that starts at city 0, visits exactly the cities in mask, and ends at j. That turns the (n-1)! orderings of brute force into 2^n · n² work.

Watch it run

The animation gives every item a lane: a is the 1, b the 2, c the 4, so one integer describes a whole subset. mask = 0 means every lane is off, the empty subset, which counts too. Add one and lane a switches on, so mask 1 is the subset holding just a. At 2 the carry turns a off and b on: counting in binary is walking the subsets. 3 has two lanes lit, a two-item subset with no recursion and no bookkeeping. 4 carries all the way up to c and clears the lower lanes. 5 is a with c, and any lane can be read with mask >> i & 1. 6 picks the other pair; union is | and intersection &, one instruction each. 7 lights every lane, the full set and the last integer below 1 << 3. Eight integers, eight subsets, in a plain loop, and each one is a hashable key for memoisation.

The code

Set operations on masks over four items:

items = ["a", "b", "c", "d"]

def members(mask):
    return [items[i] for i in range(len(items)) if mask >> i & 1]

x, y = 0b0101, 0b0011                     # a and c, a and b
print(members(x | y), members(x & y), members(x & ~y), members(x ^ y))
# ['a', 'b', 'c'] ['a'] ['c'] ['b', 'c']
print(bin(x | 1 << 3), x >> 2 & 1, bin(x).count("1"))   # 0b1101 1 2

Every subset from a counting loop, and every submask of one mask:

def subset_sums(nums):
    sums = []
    for mask in range(1 << len(nums)):    # every subset, one integer each
        sums.append(sum(nums[i] for i in range(len(nums)) if mask >> i & 1))
    return sums

print(subset_sums([3, 5, 9]))             # [0, 3, 5, 8, 9, 12, 14, 17]


def submasks(mask):
    sub, out = mask, []
    while sub:
        out.append(sub)
        sub = (sub - 1) & mask            # next smaller subset of mask
    return out + [0]

print(submasks(0b1011))                   # [11, 10, 9, 8, 3, 2, 1, 0]

Bitmask DP for the shortest round trip that visits every city once, starting and ending at city 0:

def shortest_tour(dist):
    """Bitmask DP: best[mask][j] = cheapest path from 0 visiting mask, ending at j."""
    n, INF = len(dist), float("inf")
    best = [[INF] * n for _ in range(1 << n)]
    best[1][0] = 0
    for mask in range(1 << n):
        for j in range(n):
            if best[mask][j] == INF:
                continue
            for k in range(n):
                if not mask >> k & 1:     # k not visited yet
                    nxt = mask | 1 << k
                    best[nxt][k] = min(best[nxt][k], best[mask][j] + dist[j][k])
    full = (1 << n) - 1
    return min(best[full][j] + dist[j][0] for j in range(n))

dist = [[0, 2, 9, 10], [1, 0, 6, 4], [15, 7, 0, 8], [6, 3, 12, 0]]
print(shortest_tour(dist))                # 21

Increasing masks are a valid DP order because adding a city only ever makes the mask larger. Checked with seeded random inputs against independent references: subset sums against itertools.combinations for every size, submasks against a filter over all integers, and the tour against trying every permutation of the cities:

import random
from itertools import combinations, permutations

random.seed(27)
ok = True
for _ in range(200):
    n = random.randint(0, 8)
    nums = [random.randint(-20, 20) for _ in range(n)]
    by_combinations = sorted(sum(c) for k in range(n + 1) for c in combinations(nums, k))
    ok &= sorted(subset_sums(nums)) == by_combinations
    mask = random.randrange(1 << n) if n else 0
    ok &= sorted(submasks(mask)) == [s for s in range(1 << n) if s & mask == s]
for _ in range(100):
    n = random.randint(2, 7)
    d = [[0 if i == j else random.randint(1, 50) for j in range(n)] for i in range(n)]
    brute = min(sum(d[a][b] for a, b in zip((0,) + p, p + (0,)))
                for p in permutations(range(1, n)))
    ok &= shortest_tour(d) == brute
print(ok)                                 # True

The complexity

• Set operations on one mask: O(1) for masks that fit in a machine word; Python integers grow as needed.
• Enumerating all subsets: 2^n masks; reading each subset's members adds a factor of n.
• All submasks of all masks: 3^n.
• Shortest-tour DP: O(2^n · n²) time and O(2^n · n) memory, practical up to roughly n = 20, against (n-1)! orderings for brute force.

Where it goes wrong

• Operator precedence. In Python, shifts bind tighter than & and comparisons looser, so mask & 1 << i == 0 works. In C, Java and JavaScript, == binds tighter than &, and the same line is a bug. Parenthesise anyway.
• Negative masks. ~mask in Python is negative; combine it with & against a real mask rather than using it alone.
• Using masks for large n. 2^n grows fast. Past about 25 items, enumeration is too slow whatever the representation.
• Forgetting the empty subset. Mask 0 is a subset; the submask loop above adds it explicitly because while sub stops before it.

When it shows up in interviews

Directly in "generate all subsets", where it is the iterative alternative to backtracking, and in bitmask DP problems: shortest path visiting all nodes, assigning tasks to people, partitioning into equal-sum groups. The signals, a tiny n and a question about subsets, are listed under bit tricks on the patterns cheat sheet. The single-bit operations are covered in bit masks: set, clear, toggle, and counting set bits in Brian Kernighan's trick.

How to say it in an interview

"With n this small I represent a subset as an integer: bit i is set when item i is in. Membership is mask shifted right by i, and one; union is bitwise or, intersection is bitwise and. Counting from 0 to 2 to the n minus 1 visits every subset exactly once, so enumeration is a plain loop. For the DP I use the mask as the state, cheapest cost having visited exactly this set and ending at j, and process masks in increasing order, since transitions only add bits. That is 2 to the n times n squared instead of n factorial."