Bitmask DP: Subsets, TSP, and Assignment Problems

Some problems ask you to consider every possible subset of a set. When the set is small (n <= 20), you can represent each subset as an integer where bit i indicates whether element i is in the subset. This encoding lets you use the subset as a DP state — hence the name bitmask DP.

---

1. Bit Operations Refresher

Before diving into DP, let us nail down the bit operations you will use constantly.

# Given a set of n elements, each subset is an integer in [0, 2^n)

n = 4  # elements: {0, 1, 2, 3}

# Check if element i is in the subset (mask)
def has_element(mask, i):
    return (mask >> i) & 1 == 1

# Add element i to the subset
def add_element(mask, i):
    return mask | (1 << i)

# Remove element i from the subset
def remove_element(mask, i):
    return mask & ~(1 << i)

# Toggle element i
def toggle_element(mask, i):
    return mask ^ (1 << i)

# Count elements in the subset
def count_elements(mask):
    return bin(mask).count('1')

# Iterate over all elements in the subset
def list_elements(mask):
    elements = []
    i = 0
    temp = mask
    while temp:
        if temp & 1:
            elements.append(i)
        temp >>= 1
        i += 1
    return elements

# Iterate over all subsets of a given mask
def all_subsets(mask):
    """Enumerate all subsets of mask (including mask itself and 0)."""
    sub = mask
    subsets = []
    while sub > 0:
        subsets.append(sub)
        sub = (sub - 1) & mask
    subsets.append(0)
    return subsets


# Examples
mask = 0b1011  # = 11, represents {0, 1, 3}
print(f"mask = {mask} = {bin(mask)}")
print(f"elements: {list_elements(mask)}")          # [0, 1, 3]
print(f"has element 2: {has_element(mask, 2)}")    # False
print(f"add element 2: {bin(add_element(mask, 2))}")  # 0b1111
print(f"count: {count_elements(mask)}")            # 3
print(f"subsets of 0b1010: {[bin(s) for s in all_subsets(0b1010)]}")

---

2. The Travelling Salesman Problem (TSP)

Problem: given n cities and a distance matrix dist[i][j], find the shortest route that visits every city exactly once and returns to the starting city.

This is NP-hard in general, but for small n (up to ~20), bitmask DP solves it in O(n^2 * 2^n).

State Definition

• dp[mask][i] = minimum cost to visit exactly the cities in mask, ending at city i.
• mask is a bitmask where bit j is set if city j has been visited.

Recurrence

To arrive at state (mask, i), we came from some city j that is in mask but is not i:

dp[mask][i] = min(dp[mask ^ (1 << i)][j] + dist[j][i])
              for all j in mask where j != i

Python Implementation

def tsp(dist):
    """
    Solve the Travelling Salesman Problem using bitmask DP.

    Args:
        dist: n x n distance matrix

    Returns:
        Minimum cost of a tour visiting all cities and returning to start.

    Time: O(n^2 * 2^n)
    Space: O(n * 2^n)
    """
    n = len(dist)
    INF = float('inf')
    full_mask = (1 << n) - 1

    # dp[mask][i] = min cost to reach city i having visited cities in mask
    dp = [[INF] * n for _ in range(1 << n)]

    # Start at city 0
    dp[1][0] = 0  # mask = 0b1 (only city 0 visited), at city 0

    for mask in range(1, 1 << n):
        for i in range(n):
            if dp[mask][i] == INF:
                continue
            if not (mask & (1 << i)):
                continue  # city i must be in mask

            # Try visiting an unvisited city j next
            for j in range(n):
                if mask & (1 << j):
                    continue  # already visited
                new_mask = mask | (1 << j)
                new_cost = dp[mask][i] + dist[i][j]
                if new_cost < dp[new_mask][j]:
                    dp[new_mask][j] = new_cost

    # Find minimum cost to return to city 0 after visiting all cities
    result = INF
    for i in range(n):
        if dp[full_mask][i] + dist[i][0] < result:
            result = dp[full_mask][i] + dist[i][0]

    return result


# Example: 4 cities
dist = [
    [0, 10, 15, 20],
    [10, 0, 35, 25],
    [15, 35, 0, 30],
    [20, 25, 30, 0],
]

print(tsp(dist))  # 80 (0->1->3->2->0: 10+25+30+15=80)

Reconstructing the Path

def tsp_with_path(dist):
    """TSP that also returns the optimal tour."""
    n = len(dist)
    INF = float('inf')
    full_mask = (1 << n) - 1

    dp = [[INF] * n for _ in range(1 << n)]
    parent = [[-1] * n for _ in range(1 << n)]

    dp[1][0] = 0

    for mask in range(1, 1 << n):
        for i in range(n):
            if dp[mask][i] == INF or not (mask & (1 << i)):
                continue
            for j in range(n):
                if mask & (1 << j):
                    continue
                new_mask = mask | (1 << j)
                new_cost = dp[mask][i] + dist[i][j]
                if new_cost < dp[new_mask][j]:
                    dp[new_mask][j] = new_cost
                    parent[new_mask][j] = i

    # Find the best ending city
    best_cost = INF
    last_city = -1
    for i in range(n):
        cost = dp[full_mask][i] + dist[i][0]
        if cost < best_cost:
            best_cost = cost
            last_city = i

    # Reconstruct path
    path = []
    mask = full_mask
    city = last_city
    while city != -1:
        path.append(city)
        prev = parent[mask][city]
        mask ^= (1 << city)
        city = prev

    path.reverse()
    path.append(0)  # return to start
    return best_cost, path


cost, path = tsp_with_path(dist)
print(f"Cost: {cost}, Path: {path}")

---

3. The Assignment Problem

Problem: given n people and n tasks with a cost matrix cost[i][j] (cost for person i to do task j), assign each person to exactly one task to minimise total cost.

This is solvable with the Hungarian algorithm in O(n^3), but bitmask DP gives a clean O(n * 2^n) solution for small n.

def assignment_problem(cost):
    """
    Solve the assignment problem using bitmask DP.

    Args:
        cost: n x n cost matrix where cost[i][j] = cost for person i to do task j

    Returns:
        Minimum total cost of assignment.

    Time: O(n * 2^n)
    Space: O(2^n)
    """
    n = len(cost)
    INF = float('inf')

    # dp[mask] = min cost to assign tasks in mask to the first popcount(mask) people
    dp = [INF] * (1 << n)
    dp[0] = 0  # no tasks assigned, no people used

    for mask in range(1 << n):
        if dp[mask] == INF:
            continue

        # Number of bits set = number of people already assigned
        person = bin(mask).count('1')
        if person >= n:
            continue

        # Try assigning each unassigned task to this person
        for task in range(n):
            if mask & (1 << task):
                continue  # task already assigned
            new_mask = mask | (1 << task)
            new_cost = dp[mask] + cost[person][task]
            if new_cost < dp[new_mask]:
                dp[new_mask] = new_cost

    return dp[(1 << n) - 1]


# Example
cost = [
    [9, 2, 7, 8],
    [6, 4, 3, 7],
    [5, 8, 1, 8],
    [7, 6, 9, 4],
]

print(assignment_problem(cost))  # 13 (person 0->task 1, 1->task 2? Let's see: 2+3+1+4=10? Actually 2+6+1+4=13? Let's verify.)

---

4. Counting Subsets with a Property

Problem: given an array of n numbers, count the number of subsets whose elements have XOR equal to a target value.

def count_subsets_with_xor(nums, target):
    """
    Count subsets whose XOR equals target.

    Uses bitmask to enumerate all subsets.
    Time: O(2^n), Space: O(1)
    """
    n = len(nums)
    count = 0

    for mask in range(1 << n):
        xor_val = 0
        for i in range(n):
            if mask & (1 << i):
                xor_val ^= nums[i]
        if xor_val == target:
            count += 1

    return count


print(count_subsets_with_xor([1, 2, 3, 4], 3))

For larger arrays where XOR values are bounded, use a DP approach instead:

def count_subsets_xor_dp(nums, target):
    """
    Count subsets with XOR = target using DP on XOR values.

    Time: O(n * max_xor), Space: O(max_xor)
    """
    # Find the maximum possible XOR
    max_val = 0
    for num in nums:
        max_val = max(max_val, num)

    # Upper bound for XOR
    upper = 1
    while upper <= max_val:
        upper <<= 1

    # dp[x] = number of subsets with XOR = x
    dp = [0] * upper
    dp[0] = 1  # empty subset has XOR = 0

    for num in nums:
        new_dp = dp[:]
        for x in range(upper):
            new_dp[x ^ num] += dp[x]
        dp = new_dp

    return dp[target] if target < upper else 0


print(count_subsets_xor_dp([1, 2, 3, 4], 3))

---

5. Minimum Cost to Visit All Nodes in a Graph

A variant of TSP for general graphs (not necessarily complete). You can revisit nodes, and you want the shortest walk that visits every node at least once.

from collections import deque

def shortest_path_visiting_all_nodes(graph):
    """
    LeetCode 847: Shortest Path Visiting All Nodes.

    BFS with state = (current_node, visited_mask).

    Time: O(n * 2^n)
    Space: O(n * 2^n)
    """
    n = len(graph)
    full_mask = (1 << n) - 1

    # BFS: state = (node, mask), value = distance
    queue = deque()
    visited = set()

    # Start BFS from every node
    for i in range(n):
        state = (i, 1 << i)
        queue.append((i, 1 << i, 0))
        visited.add(state)

    while queue:
        node, mask, dist = queue.popleft()

        if mask == full_mask:
            return dist

        for neighbor in graph[node]:
            new_mask = mask | (1 << neighbor)
            state = (neighbor, new_mask)
            if state not in visited:
                visited.add(state)
                queue.append((neighbor, new_mask, dist + 1))

    return -1  # should never reach here for connected graph


# Example
graph = [[1, 2, 3], [0], [0], [0]]
print(shortest_path_visiting_all_nodes(graph))  # 4

---

6. Partition Into K Equal Sum Subsets

Problem: given an array nums and integer k, determine if the array can be partitioned into k subsets with equal sum.

def can_partition_k_subsets(nums, k):
    """
    LeetCode 698: Partition to K Equal Sum Subsets.

    dp[mask] = number of complete groups formed when
    elements in mask have been used, plus the remainder.

    Time: O(n * 2^n)
    Space: O(2^n)
    """
    total = sum(nums)
    if total % k != 0:
        return False

    target = total // k
    n = len(nums)
    nums.sort(reverse=True)  # optimization: try large nums first

    if nums[0] > target:
        return False

    # dp[mask] = the current sum modulo target for the elements used so far
    # -1 means this state is unreachable
    dp = [-1] * (1 << n)
    dp[0] = 0

    for mask in range(1 << n):
        if dp[mask] == -1:
            continue

        for i in range(n):
            if mask & (1 << i):
                continue  # already used

            # Can we add nums[i] to the current group?
            if dp[mask] + nums[i] <= target:
                new_mask = mask | (1 << i)
                dp[new_mask] = (dp[mask] + nums[i]) % target

    return dp[(1 << n) - 1] == 0


print(can_partition_k_subsets([4, 3, 2, 3, 5, 2, 1], 4))  # True (target=5)

---

7. Bitmask DP Over Subsets of Subsets

Sometimes you need to iterate over all subsets of a given subset. The key trick:

def iterate_subsets(mask):
    """
    Iterate over all non-empty subsets of mask in O(2^popcount(mask)) time.
    This is Gosper's hack for subset enumeration.
    """
    sub = mask
    while sub > 0:
        print(bin(sub))
        sub = (sub - 1) & mask


# Total work across all masks: sum over all masks of 2^popcount(mask) = 3^n
# This is because each element is in one of 3 states:
# not in mask, in mask but not in sub, in mask and in sub
iterate_subsets(0b1011)
# Output: 0b1011, 0b1010, 0b1001, 0b1000, 0b11, 0b10, 0b1

Application: Minimum Steiner Tree

def min_cost_steiner_tree(n, edges, terminals):
    """
    Find the minimum cost subgraph connecting all terminal nodes.
    Simplified version using bitmask DP over terminal subsets.
    """
    INF = float('inf')
    k = len(terminals)

    # Build adjacency list
    adj = [[] for _ in range(n)]
    for u, v, w in edges:
        adj[u].append((v, w))
        adj[v].append((u, w))

    # Map terminal indices
    term_idx = {t: i for i, t in enumerate(terminals)}

    # dp[mask][v] = min cost tree connecting terminals in mask, rooted at v
    dp = [[INF] * n for _ in range(1 << k)]

    # Base case: single terminals
    for i, t in enumerate(terminals):
        dp[1 << i][t] = 0

    # Fill DP
    from heapq import heappush, heappop

    for mask in range(1, 1 << k):
        # Merge subsets
        sub = (mask - 1) & mask
        while sub > 0:
            comp = mask ^ sub
            if sub < comp:  # avoid double counting
                sub = (sub - 1) & mask
                continue
            for v in range(n):
                val = dp[sub][v] + dp[comp][v]
                if val < dp[mask][v]:
                    dp[mask][v] = val
            sub = (sub - 1) & mask

        # Dijkstra to propagate within this mask
        heap = []
        for v in range(n):
            if dp[mask][v] < INF:
                heappush(heap, (dp[mask][v], v))

        while heap:
            cost, u = heappop(heap)
            if cost > dp[mask][u]:
                continue
            for v, w in adj[u]:
                if cost + w < dp[mask][v]:
                    dp[mask][v] = cost + w
                    heappush(heap, (dp[mask][v], v))

    full = (1 << k) - 1
    return min(dp[full])

---

8. When to Use Bitmask DP

Use bitmask DP when:

• The set size n is small (n <= 20 for O(n * 2^n), n <= 15 for O(n^2 * 2^n))
• You need to track which elements have been “used” or “visited”
• The problem involves permutations, assignments, or subsets

Do not use bitmask DP when:

• n is large (use greedy, flow, or other polynomial algorithms)
• The problem has a polynomial-time solution (check first!)
• Only subset sums matter (use knapsack DP instead)

Complexity Comparison

Approach	Time	Max n
Brute force (all permutations)	O(n!)	~10
Bitmask DP	O(n^2 * 2^n) or O(n * 2^n)	~20
Subset enumeration (3^n)	O(3^n)	~15
Polynomial algorithms	O(n^2) or O(n^3)	thousands+

---

9. Practice Problems

Problem	Platform	Key Technique
Shortest Path Visiting All Nodes (LC 847)	LeetCode	BFS + bitmask
Partition to K Equal Sum Subsets (LC 698)	LeetCode	Bitmask DP
Parallel Courses II (LC 1494)	LeetCode	Bitmask + subsets
Can I Win (LC 464)	LeetCode	Game theory + bitmask
Travelling Salesman (CSES)	CSES	Classic TSP
Hamiltonian Flights (CSES)	CSES	Count Hamiltonian paths
Maximum Students Taking Exam (LC 1349)	LeetCode	Bitmask per row
Number of Ways to Wear Different Hats (LC 1434)	LeetCode	Assignment DP

---

10. Big-O Summary

Algorithm	Time	Space
TSP (bitmask DP)	O(n^2 * 2^n)	O(n * 2^n)
Assignment problem	O(n * 2^n)	O(2^n)
Subset enumeration	O(2^n)	O(1)
Subset of subsets	O(3^n) total	O(2^n)
BFS + bitmask	O(n * 2^n)	O(n * 2^n)

Bitmask DP is the bridge between brute force and polynomial algorithms. When n is too large for brute force but too small or too structured for polynomial algorithms, bitmask DP hits the sweet spot. The key is recognising that the “which elements have been used” state can be compressed into a single integer.