Bitmask DP

Problem

The questions look like search problems: give each worker one job at the lowest total cost, visit every city once on the shortest tour, split numbers into k groups of equal sum, cover a list of skills with the smallest team. Each choice uses up an item, and later choices are only limited by what is used up. The constraint that gives it away is a tiny n, usually 12 to 20, which is too small for a polynomial algorithm to be the point and too big for n! to finish.

Intuition

Three workers, A, B and C, must each take one of jobs 0, 1 and 2. Trying every order means 3! = 6 assignments, and for 20 workers, 20! of them. Look at what those orders have in common. Once A and B have taken jobs 0 and 1, between them, C must take job 2, whichever of the two orders put them there. Both partial assignments face the same future, so only the cheaper one can be part of the best answer. Keep that one and throw the other away.

That is the whole idea: for every set of jobs already given out, remember only the cheapest way to give them out. A set of jobs is a row of yes/no answers, and a row of yes/no answers is an integer written in binary. Set {0, 1} is 011 = 3, set {0, 2} is 101 = 5. The DP table is then a plain list indexed by those integers, and adding job j to a set is one bit operation.

Greedy reasoning does not survive this. If each worker in turn takes their cheapest free job, A grabs job 1 for 2, and the total is 10. The best total is 9, because the DP also keeps the partial assignments that look worse early on.

Green masks lie on the cheapest path from 000 to 111, and the labels give the cost each step adds. Every mask keeps one number, however many paths reach it.

Mechanics

The bit idioms

With n items numbered 0 to n − 1, the set is an integer whose bit j is 1 when item j is in it. Test membership with mask >> j & 1, add an item with mask | 1 << j, remove one with mask & ~(1 << j), and count members with mask.bit_count() (Python 3.10+; bin(mask).count("1") before that). The empty set is 0 and the full set is (1 << n) − 1. These are the tricks from the bit manipulation page, used as a data structure.

State and transition

For the assignment problem, dp[mask] is the lowest cost of giving the jobs in mask to the first popcount(mask) workers. Worker i = popcount(mask) is next, and for each free job j, dp[mask | 1 << j] may improve to dp[mask] + cost[i][j]. Start from dp[0] = 0 and read the answer at the full mask. Adding a bit always makes the integer bigger, so looping over masks in increasing numeric order handles every mask after all the masks it is built from.

Why it is correct. Take any best complete assignment and stop it after k workers; the jobs they hold form some mask. If a cheaper way to fill that same mask existed, swapping it in would leave the remaining workers exactly the same jobs to choose from, and give a cheaper complete assignment, which is impossible. So every prefix of a best answer is a best answer for its own mask, and the DP, which keeps the best answer for every mask, cannot lose it. The argument only works because the rest of the problem depends on the set and nothing else. That is the question to ask before using this pattern: is the set enough to describe the future?

When the set is not enough

On a tour, the next leg's cost depends on where you are standing, so the set of visited cities is not enough. Add the last city to the state: dp[mask][last] is the shortest path that starts at city 0, visits exactly the cities in mask and ends at last. It extends to dp[mask | 1 << c][c] through dist[last][c]. This is the Held–Karp algorithm (Bellman found it the same year), and it is still the best known exact method for the general travelling salesman problem in the worst case. A tour is a cycle, so its start can be fixed at city 0 without losing any tour, which saves a factor of n.

Worked example

Assignment. The cost of each worker doing each job:

Every mask, in order of size, with the step that gave it its value:

Mask 3 (011) shows the pruning at work: A on job 0 then B on job 1 costs 13, while A on job 1 then B on job 0 costs 8, so only 8 is kept. Following the best last steps back from 111 gives A → job 1, B → job 0, C → job 2, for a total of 9. The greedy order costs 10.

Held–Karp on four cities. The distances:

Every reachable state, starting from city 0:

Close each full path by returning to city 0:

The shortest tour is 0 → 2 → 3 → 1 → 0, length 80. Two endings tie at 80 because the distances are symmetric: each is the other tour driven backwards.

City 0 is the fixed start. Green edges form the shortest tour; dashed edges are the others. Each box on the right is one DP state the tour passes through, and its bits read city 3, 2, 1, 0.

Implementation

import math


def min_assignment_cost(cost):
    n = len(cost)
    dp = [math.inf] * (1 << n)
    dp[0] = 0
    for mask in range(1 << n):
        if dp[mask] == math.inf:
            continue
        worker = mask.bit_count()
        if worker == n:
            continue
        for job in range(n):
            if not mask >> job & 1:
                bigger = mask | 1 << job
                dp[bigger] = min(dp[bigger], dp[mask] + cost[worker][job])
    return dp[(1 << n) - 1]

Variations and edge cases

Top-down. The same recurrence with a cache is often quicker to write, and visits only the masks it needs:

from functools import cache

def min_assignment_cost(cost):
    n = len(cost)

    @cache
    def best(mask):
        worker = mask.bit_count()
        if worker == n:
            return 0
        return min(cost[worker][job] + best(mask | 1 << job)
                   for job in range(n) if not mask >> job & 1)

    return best(0)

Held–Karp. Two dimensions, the set and the last city:

import math


def shortest_tour(dist):
    n = len(dist)
    dp = [[math.inf] * n for _ in range(1 << n)]
    dp[1][0] = 0
    for mask in range(1, 1 << n, 2):
        for last in range(n):
            here = dp[mask][last]
            if here == math.inf:
                continue
            for city in range(1, n):
                if not mask >> city & 1:
                    bigger = mask | 1 << city
                    dp[bigger][city] = min(dp[bigger][city], here + dist[last][city])
    full = (1 << n) - 1
    return min(dp[full][last] + dist[last][0] for last in range(1, n))

Groups of equal sum. Fill groups one at a time; fill[mask] is how full the open group is after using the items in mask, and an item may only join if it fits. The answer is yes when the full mask is reachable:

def can_partition_k(nums, k):
    total = sum(nums)
    if total % k:
        return False
    side, n = total // k, len(nums)
    fill = [-1] * (1 << n)
    fill[0] = 0
    for mask in range(1 << n):
        if fill[mask] < 0:
            continue
        for i in range(n):
            if not mask >> i & 1 and fill[mask] + nums[i] <= side:
                fill[mask | 1 << i] = (fill[mask] + nums[i]) % side
    return fill[-1] == 0

Every submask of a mask. Some recurrences split a set into two parts. The step sub = (sub − 1) & mask walks through every submask from largest to 0, skipping the integers that are not submasks. For mask 11 (1011) it visits 1011, 1010, 1001, 1000, 0011, 0010, 0001, 0000: 8 of them, one per subset of its 3 set bits. Over all masks of n bits the total is 3ⁿ, since each item is outside the mask, inside it but not the submask, or in both.

sub = mask
while True:
    visit(sub)
    if sub == 0:
        break
    sub = (sub - 1) & mask

• Unreachable states. Start every entry at infinity (or −1) and skip it when extending. Otherwise an unreachable state passes its default value along as if it were real.
• Operator precedence. In Python, << and >> bind tighter than & and |, so mask & 1 << j means what it says. Write the brackets anyway if they help you read it.
• Memory. dp[mask][last] with n = 20 is about 21 million entries. In Python that is too slow and too big; keep n at 16 or so for two dimensions, or store each row as an array.
• Fixing a start. For a cycle, fixing city 0 is safe. For a path that may start anywhere, seed dp[1 << i][i] = 0 for every i instead.

Complexity and tradeoffs

The assignment DP has 2ⁿ masks and tries n jobs from each: O(n · 2ⁿ) time and O(2ⁿ) space. Held–Karp has n · 2ⁿ states and tries n next cities from each: O(n² · 2ⁿ) time and O(n · 2ⁿ) space. The brute-force baseline tries every order, n!:

• Where the limit sits. At about 10⁷ to 10⁸ simple steps per second, n · 2ⁿ is comfortable up to n = 20, n² 2ⁿ up to about 16, and 3ⁿ up to about 15. Constraints that stop just there are a hint.
• Against backtracking. For yes/no questions like equal-sum groups, backtracking with good pruning often finishes much faster in practice; the DP's advantage is a guaranteed bound.
• Against polynomial algorithms. Assignment is solvable in O(n³) by the Hungarian method, so for large n bitmask DP is the wrong tool. Travelling salesman is NP-hard, so nothing polynomial is known; for large n, real solvers use branch and bound or heuristics.
• When the set is not enough. If the future depends on the order beyond a last element or two, the state grows past 2ⁿ · n and the pattern no longer pays.

Problems

No problem on the four lists has this page as its home, which is why it sits under Beyond the interview.

Also uses this pattern

These problems have a different home page, but this idea is part of their solution.

• 473. Matchsticks to Square Medium
• 698. Partition to K Equal Sum Subsets Medium

For practice beyond the lists, in rising difficulty:

• 526. Beautiful Arrangement Medium
  Place numbers into positions 1, 2, 3 in order, so the size of the used set tells you which position comes next; count ways per mask.
• 464. Can I Win Medium
  The set of numbers taken fixes the running total, so a game state is just a mask, and a player wins if some move leaves the other a losing mask.
• 1947. Maximum Compatibility Score Sum Medium
  This page's assignment DP with a maximum instead of a minimum, after scoring every student and mentor pair.
• 1879. Minimum XOR Sum of Two Arrays Hard
  Pairing two arrays is an assignment where the cost of a pair is the XOR of its values.
• 847. Shortest Path Visiting All Nodes Hard
  Run BFS over pairs of current node and visited mask, starting from every node at once.
• 1125. Smallest Sufficient Team Hard
  Here the mask is the set of skills covered, not people used; keep the smallest team for each skill mask.
• 1494. Parallel Courses II Hard
  From each mask of finished courses, try every submask of the available ones that fits in a term, which is where 3ⁿ comes from.
• 943. Find the Shortest Superstring Hard
  Held–Karp in disguise: the cost of going from word a to word b is how much of b does not overlap a's end.

Connections

Bitmask DP is DP whose state happens to be a set, stored through bit manipulation. It is what permutation backtracking becomes once you notice that many orders reach the same set: the cache merges them. Its yes/no cousins, Matchsticks to Square and Partition to K Equal Sum Subsets, live on the constraint backtracking page, where pruning usually beats the DP in practice. BFS over (node, mask) pairs is graph traversal on a state graph, and visiting every node once is the Hamiltonian cousin of the Eulerian path, which uses every edge once and is far easier.

Check yourself

References

• Held, M. and Karp, R. M., "A dynamic programming approach to sequencing problems", Journal of the Society for Industrial and Applied Mathematics 10(1):196–210, 1962.
• Bellman, R., "Dynamic programming treatment of the travelling salesman problem", Journal of the ACM 9(1):61–63, 1962.
• Kuhn, H. W., "The Hungarian method for the assignment problem", Naval Research Logistics Quarterly 2:83–97, 1955.
• Cormen, Leiserson, Rivest and Stein, Introduction to Algorithms, 3rd edition, chapter 15, "Dynamic Programming", and section 34.5.4 on the travelling-salesman problem.
• Skiena, S., The Algorithm Design Manual, the catalog entry "Traveling Salesman Problem".
• Sedgewick and Wayne, Algorithms, 4th edition, chapter 6, "Context", on intractability.
• Knuth, D. E., The Art of Computer Programming, volume 4A, section 7.1.3, "Bitwise Tricks and Techniques".
• The NeetCode roadmap and the Blind 75, Grind 169, NeetCode 150 and NeetCode 250 lists, used as the index of topics and problems.