Bitmask DP in Brief
24 min · Intermediate
Need help? Ask a tutor on WhatsApp (opens in a new tab)When there are at most about 20 things, a set of them fits in one integer, one bit per thing. A table indexed by that integer can then solve assignment and ordering problems.
A set stored in one integer
A bitmask is an integer used as a set: bit i is 1 when item i is in the set. Adding, removing and testing members uses the operators |, & and shifts such as 1 << i.
Number the items 0 to n - 1. Bit i of an integer is 1 when item i is in the set. Python integers have no size limit, and bit_count() counts the members.
mask = 0mask |= 1 << 2 # add item 2mask |= 1 << 4 # add item 4print(bin(mask), mask.bit_count())print(bool(mask >> 4 & 1), bool(mask >> 3 & 1))mask &= ~(1 << 2) # remove item 2print(bin(mask))Example 1 · Sets of skills as masks
python_skill, sql_skill, git_skill = 1, 2, 4needed = python_skill | sql_skillalice = python_skill | git_skillbob = python_skill | sql_skill | git_skillprint(alice & needed == needed, bob & needed == needed)print(bin(alice & bob), (alice | bob).bit_count())False True 0b101 3
What this code does
&is the intersection and|is the union. A set contains another whena & b == b.- Alice lacks the SQL skill, so she fails the test. Bob passes.
Example 2 · Count the members of every mask
print([m.bit_count() for m in range(8)])print(2 ** 20, (1 << 20) - 1)[0, 1, 1, 2, 1, 2, 2, 3] 1048576 1048575
What this code does
- The list shows how many items each 3-item set holds. Mask 7 holds all three.
- Twenty items give about a million masks, so the table is large but fine to hold in memory.
A set of n items has 2 to the power n possible masks, numbered 0 to 2^n - 1. That is why a plain list can be indexed by a set.
Assign jobs: dp over sets
Give n workers n jobs, one each, at the least total cost. dp[mask] is the cheapest way to give the first popcount(mask) workers exactly the jobs in mask. The next worker takes any job not yet in the mask.
def min_assignment(cost): n = len(cost) INF = float("inf") dp = [INF] * (1 << n) dp[0] = 0 for mask in range(1 << n): worker = mask.bit_count() if worker == n: continue for job in range(n): if not mask >> job & 1: nxt = mask | 1 << job dp[nxt] = min(dp[nxt], dp[mask] + cost[worker][job]) return dp[-1]Example 3 · Cheapest job assignment
def min_assignment(cost): n = len(cost) INF = float("inf") dp = [INF] * (1 << n) dp[0] = 0 for mask in range(1 << n): worker = mask.bit_count() if worker == n: continue for job in range(n): if not mask >> job & 1: nxt = mask | 1 << job dp[nxt] = min(dp[nxt], dp[mask] + cost[worker][job]) return dp[-1]
cost = [[9, 2, 7], [6, 4, 3], [5, 8, 1]]print(min_assignment(cost))9
What this code does
- Rows are workers and columns are jobs. The best choice is worker 0 on job 1 (2), worker 1 on job 0 (6) and worker 2 on job 2 (1), a total of 9.
- dp[mask] is built from smaller masks, since a mask with one more job is always a bigger integer.
Example 4 · Check against brute force
from itertools import permutations
cost = [[9, 2, 7], [6, 4, 3], [5, 8, 1]]best = min(sum(cost[w][job] for w, job in enumerate(order)) for order in permutations(range(3)))print(best)9
What this code does
permutationstries every way to give the jobs, 6 for three workers. It matches the table answer.- Brute force is fine for 8 workers and hopeless for 20. The bitmask table is the next step up.
There are 2 to the power n masks and n choices for each, so the cost is O(2^n x n) time and O(2^n) space. Trying all orders by brute force is n factorial.
A tour of cities: state is (visited, last)
Visit every city once and return to the start at the least distance. Brute force tries n factorial orders. The state is the set of visited cities and the city you stand in: 2^n x n values, each with n moves.
def tour(dist): n = len(dist) INF = float("inf") dp = [[INF] * n for _ in range(1 << n)] dp[1][0] = 0 for mask in range(1 << n): for last in range(n): if dp[mask][last] == INF: continue for nxt in range(n): if not mask >> nxt & 1: new = mask | 1 << nxt dp[new][nxt] = min(dp[new][nxt], dp[mask][last] + dist[last][nxt]) return min(dp[-1][last] + dist[last][0] for last in range(n))Example 5 · The shortest tour of four cities
def tour(dist): n = len(dist) INF = float("inf") dp = [[INF] * n for _ in range(1 << n)] dp[1][0] = 0 for mask in range(1 << n): for last in range(n): if dp[mask][last] == INF: continue for nxt in range(n): if not mask >> nxt & 1: new = mask | 1 << nxt dp[new][nxt] = min(dp[new][nxt], dp[mask][last] + dist[last][nxt]) return min(dp[-1][last] + dist[last][0] for last in range(n))
dist = [ [0, 10, 15, 20], [10, 0, 35, 25], [15, 35, 0, 30], [20, 25, 30, 0],]print(tour(dist))80
What this code does
- The best round trip is 0 to 1 to 3 to 2 and back to 0, with 10 + 25 + 30 + 15 = 80.
- dp[mask][last] is the shortest way to visit the cities in mask and finish at last.
Mask 1 means only city 0 is visited. The answer adds the way home from every last city. The cost is O(2^n x n squared) time and O(2^n x n) space.
Submasks and the size limit
The subsets of a set are the masks hidden inside it, and the loop below visits each non-empty one once. Bitmask DP costs at least 2 to the power n, so in Python it suits n up to about 15 to 18.
mask = 0b1011sub = maskwhile sub: print(bin(sub)) sub = (sub - 1) & maskExample 6 · Visit every submask
mask = 0b1011sub = maskfound = []while sub: found.append(bin(sub)) sub = (sub - 1) & maskprint(found)['0b1011', '0b1010', '0b1001', '0b1000', '0b11', '0b10', '0b1']
What this code does
- The mask has three bits, so it has 2^3 - 1 = 7 non-empty submasks.
- They appear from the largest to the smallest, and each exactly once.
Subtracting 1 and masking steps to the next smaller submask. Looping over all submasks of all masks costs 3 to the power n, so keep n small.
Next step
Check your understanding
11 quick questions