L44 · DP XII — Bitmask DP

Series: LeetCode — From Basics to Interview-Ready · Session 44 / 65 · Phase 2 · Dynamic programming

Why this session exists

Bitmask DP is only viable for n at most about 20, and the constraint tells you. That ties directly back to L01, where the whole argument was that constraints are the problem setter speaking to you in a compressed language. Nowhere is that clearer than here.

Think about what n <= 20 actually rules out. It is far too small for the constraint to be about performance in any ordinary sense — a linear or quadratic algorithm on twenty elements finishes instantly, so why would a setter bother restricting it? The only reason to cap n that low is that the intended solution is exponential in n. And 2^20 is about a million, which is exactly the size of a comfortable DP table. The constraint is not a limitation; it is a hint that arrived before the problem did.

The second thing this session buys you is fluency. Bitmask problems are not conceptually harder than other DP — the state is a subset, the transitions add one element — but the syntax is unfamiliar enough that people lose minutes fumbling with shifts and ands. That fumbling is what actually costs you in an interview, not the insight. So a chunk of this session is drilling the four or five idioms until they are automatic.

The third thing is knowing when not to reach for it. 2^n grows viciously. At n = 25 you are at 33 million states, and if each state does O(n) work you are done for. Recognising the ceiling is as valuable as recognising the pattern.

Blank-file warm-up

Five minutes, empty file, no notes. Write from memory:

1. for mask in range(1 << n) — the iteration over all subsets.
2. Test whether element i is in mask.
3. Add element i to mask, and remove it.
4. Count the set bits in mask, both with a builtin and with the n & (n-1) loop.

If any of these takes more than a few seconds to recall, that is the finding, and drilling them is more valuable than solving another problem today.

Pattern anatomy

The shape that summons this pattern: a set of at most about twenty items, where the answer depends on which subset you have already used, and order within the subset does or does not matter but membership definitely does.

The invariant: dp[mask] is the best value achievable having exactly the elements of mask already handled. Bit i of mask is 1 if element i has been used. Every subset gets exactly one integer, and every integer below 2^n is a valid subset — the correspondence is total, which is why the encoding works at all.

Here are the idioms. Learn them as vocabulary, not as puzzles:

n = 5
 
full = (1 << n) - 1          # 0b11111 — all elements present
empty = 0
 
# membership
if mask & (1 << i):          # is element i in the set?
    ...
 
# insert / remove / toggle
mask_with_i    = mask | (1 << i)
mask_without_i = mask & ~(1 << i)
mask_toggled   = mask ^ (1 << i)
 
# size of the subset
size = bin(mask).count("1")  # or mask.bit_count() on Python 3.10+
 
# lowest set bit, and clearing it
low   = mask & (-mask)       # isolates the lowest set bit as a power of two
mask2 = mask & (mask - 1)    # clears the lowest set bit
 
# iterate over every element present in mask
m = mask
while m:
    b = m & (-m)
    i = b.bit_length() - 1
    m ^= b
    ...

And the standard forward-transition skeleton — iterate masks in increasing order, and from each reachable mask push to every mask that adds one more element:

def bitmask_dp(n, cost):
    INF = float("inf")
    dp = [INF] * (1 << n)
    dp[0] = 0
 
    for mask in range(1 << n):
        if dp[mask] == INF:
            continue                       # unreachable, skip
        k = bin(mask).count("1")           # how many placed so far
        for i in range(n):
            if mask & (1 << i):
                continue                   # already used
            nxt = mask | (1 << i)
            dp[nxt] = min(dp[nxt], dp[mask] + cost(k, i))
    return dp[(1 << n) - 1]

Increasing numeric order is a valid fill order, and this is worth understanding rather than accepting: adding a bit strictly increases the integer value, so mask | (1 << i) is always numerically greater than mask. Every dependency therefore sits at a lower index and is already final.

The cue

You are looking at bitmask DP when the statement contains these tells:

1. n at most 16, 20, or 22. The single loudest tell in this series. Twenty is the canonical cap because 2^20 is about a million and 2^20 × 20 is about twenty million, which is the outer limit of comfortable. If you see n <= 20 and the problem is not trivially linear, think subsets immediately.
2. "Assign every item to some group" or "visit all nodes" or "partition into k subsets". These are all questions about which elements have been consumed, and consumption order is either irrelevant or captured by a second small dimension.
3. A permutation flavour with n small. Counting or optimising over orderings of at most twenty items — 20! is impossible but 2^20 × 20 is fine, because the DP collapses all orderings that use the same set into one state. That collapse is the entire win.
4. "Shortest path visiting all nodes" and anything else that smells of travelling-salesman. TSP is the archetype: state is (mask of visited, current node).
5. You reached for backtracking and the search tree is n!. If pruning is not obviously enough and n is small, the fix is to memoise on the set rather than the sequence.

Tell 3 is the conceptual heart. n! orderings but only 2^n sets means that if the future only depends on the set consumed rather than the order, you have collapsed a factorial into an exponential — an enormous saving that people often miss.

Guided solve

Partition to K Equal Sum Subsets — given an array nums and an integer k, decide whether the array can be split into k subsets with equal sums.

Start with the cheap rejections, because they cost two lines and eliminate most invalid inputs:

• If sum(nums) is not divisible by k, return False immediately.
• Let target = sum(nums) // k. If any single element exceeds target, return False — that element can never fit anywhere.

Now the framing. The naive approach is backtracking: try to fill bucket one, then bucket two, and so on. It works with aggressive pruning but the search tree is large and the pruning rules are fiddly to get right under pressure.

The bitmask reframing is cleaner. Do not track k separate buckets. Instead, walk through the elements in one long sequence, filling one bucket at a time, and let the running sum modulo target tell you where the bucket boundaries fall. The key realisation: if mask tells you which elements are used, then sum(mask) is determined, so the amount currently in the partially-filled bucket is sum(mask) % target — you do not need to store it. The mask alone is a sufficient state.

def canPartitionKSubsets(nums, k):
    total = sum(nums)
    if total % k:
        return False
    target = total // k
    if max(nums) > target:
        return False
 
    n = len(nums)
    nums.sort(reverse=True)                 # big first: fails fast, prunes hard
 
    full = (1 << n) - 1
    # subset_sum[mask] = sum of the chosen elements
    subset_sum = [0] * (1 << n)
    for mask in range(1, 1 << n):
        low = mask & (-mask)
        i = low.bit_length() - 1
        subset_sum[mask] = subset_sum[mask ^ low] + nums[i]
 
    dp = [False] * (1 << n)                 # dp[mask] = mask is a valid prefix state
    dp[0] = True
 
    for mask in range(1 << n):
        if not dp[mask]:
            continue
        used = subset_sum[mask] % target    # how full the current bucket is
        for i in range(n):
            if mask & (1 << i):
                continue
            if used + nums[i] > target:
                continue                    # would overflow the current bucket
            dp[mask | (1 << i)] = True
    return dp[full]

Two things carry the correctness. First, subset_sum[mask] % target is exactly the fill level of the in-progress bucket, because every completed bucket contributed exactly target and vanishes under the modulo. Second, the overflow check used + nums[i] > target is what forces buckets to be exactly full before a new one opens — a bucket can never be skipped while under target, because no element could then be added and the state would be a dead end.

Sorting descending is a real optimisation, not decoration. Large elements have the fewest valid placements, so trying them first prunes the search space early. On the flip side, note that the DP as written visits all 2^n masks regardless, so the sort helps by making the overflow check reject more transitions rather than by pruning states outright.

Solo timed

Fifteen minutes each, timer running, no editorial until it fires.

• Shortest Path Visiting All Nodes — the state is (mask, current_node), so the table has 2^n × n entries. This is a BFS over that state space rather than a fill-order DP, because edge weights are uniform. Think about what the start states are — there is more than one.

If that lands early, attempt Number of Ways to Wear Different Hats or Maximum Students Taking Exam. The latter uses a mask per row rather than a mask over all elements, which is a genuinely different flavour worth seeing.

Common failure modes

Operator precedence on the membership test. In Python, & binds looser than ==, so mask & 1 << i == 0 parses in a way you did not intend. Always parenthesise: (mask & (1 << i)) == 0. This bug is silent and infuriating.

Iterating masks in the wrong direction. Forward transitions require increasing order; if you write backward transitions that read dp[mask ^ (1 << i)], increasing order is still correct, but mixing the two styles in one loop is not. Pick push or pull and stay there.

Recomputing popcount or subset sums inside the inner loop. Turns a tight solution into a timeout. Precompute both in a single linear-in-masks pass using the lowest-set-bit recurrence.

Not skipping unreachable states. Without the if dp[mask] is unreachable: continue guard you do n units of work on states that can never contribute. On a million masks that is twenty million wasted operations for nothing.

Assuming n is small when it is the array length of something else. Read carefully which quantity the constraint bounds. nums.length <= 20 is your green light; nums[i] <= 20 is not.

Complexity

Time: O(2^n · n) for the standard skeleton — 2^n masks, each considering at most n elements to add, with constant work per consideration once popcounts and subset sums are precomputed. At n = 20 that is roughly twenty million operations, which is seconds in Python and milliseconds in C++.

Space: O(2^n) for the DP array, plus another O(2^n) if you precompute subset sums. At n = 20 a Python list of a million booleans is fine; a list of a million floats is heavier but still workable. At n = 24 you are at sixteen million entries and memory becomes the binding constraint before time does.

The TSP variant is O(2^n · n²) because the state is (mask, last_node) — 2^n · n states, each trying n next nodes. At n = 12 that is about two hundred thousand transitions; at n = 20 it is four hundred million and out of reach for Python. Compute this bound before you start writing.

Spaced queue

Re-solve whatever is due before starting anything new today. Status ladder:

• cold — solved unaided, first attempt clean → next review in 60 days
• warm — solved but slowly or with a stumble → 21 days
• hint — needed a nudge to get started → 7 days
• failed — could not produce a working solution → 2 days, then 7 days

Add a separate queue entry for the bit idioms themselves, not just for the problems. Fluency with 1 << i and mask & (-mask) decays independently of whether you remember how Partition to K Equal Sum Subsets works.