Week 18: Advanced Dynamic Programming — Bitmask & Digit DP

1. Bitmask DP: The Traveling Salesman Problem

The Traveling Salesman Problem asks for the minimum-cost route that visits every city exactly once and returns to the start. Trying every permutation of cities is O(n!) — for n = 15 that's already over a trillion routes. The key realization: the cost of the rest of a route only depends on the current city and which set of cities has already been visited, not on the specific order they were visited in. That set has 2ⁿ possible values, small enough to index directly — represent it as an n-bit integer, where bit i means "city i has been visited."

#include <vector>
#include <climits>
using namespace std;

// dist[i][j] = cost to travel directly from city i to city j
// dp[mask][i] = min cost to have visited exactly the cities in `mask`, ending at city i
// O(2^n * n^2) time, O(2^n * n) space
int tsp(vector<vector<int>>& dist) {
    int n = dist.size();
    int fullMask = (1 << n) - 1;
    vector<vector<int>> dp(1 << n, vector<int>(n, INT_MAX));

    dp[1][0] = 0;   // start at city 0, having visited only city 0 (bit 0 set)

    for (int mask = 1; mask <= fullMask; mask++) {
        for (int last = 0; last < n; last++) {
            if (!(mask & (1 << last))) continue;         // last must be in this mask
            if (dp[mask][last] == INT_MAX) continue;      // unreachable state

            for (int next = 0; next < n; next++) {
                if (mask & (1 << next)) continue;         // next already visited
                int newMask = mask | (1 << next);
                int newCost = dp[mask][last] + dist[last][next];
                if (newCost < dp[newMask][next]) {
                    dp[newMask][next] = newCost;
                }
            }
        }
    }

    // Close the tour: return from every possible last city back to city 0
    int best = INT_MAX;
    for (int last = 1; last < n; last++) {
        if (dp[fullMask][last] != INT_MAX) {
            best = min(best, dp[fullMask][last] + dist[last][0]);
        }
    }
    return best;
}

dp[mask][last] means "the minimum cost of a partial route that has visited exactly the cities marked in mask, currently sitting at city last." There are 2ⁿ masks and n possible last cities, so the table has O(2ⁿ · n) entries, and filling each one considers up to n transitions — giving O(2ⁿ · n²) overall. For n = 15, that's roughly 7 million operations instead of a trillion-plus permutations — the difference between infeasible and instant.

2. Bitmask DP: The Assignment Problem

The assignment problem: given n workers and n tasks, with cost[i][j] for assigning worker i to task j, find the minimum-cost way to assign every worker to exactly one task. This is the same bitmask shape as TSP, but the mask tracks "which tasks have been assigned" while workers are processed one at a time in a fixed order — since every worker gets assigned exactly once, you don't need the mask to also track which workers are done.

#include <vector>
#include <climits>
using namespace std;

// dp[mask] = min cost to assign the first popcount(mask) workers to the tasks in `mask`
// O(2^n * n) time
int minAssignmentCost(vector<vector<int>>& cost) {
    int n = cost.size();
    int fullMask = (1 << n) - 1;
    vector<int> dp(1 << n, INT_MAX);
    dp[0] = 0;   // no workers assigned yet, no tasks used

    for (int mask = 0; mask < fullMask; mask++) {
        if (dp[mask] == INT_MAX) continue;
        int worker = __builtin_popcount(mask);   // the next worker to assign
        for (int task = 0; task < n; task++) {
            if (mask & (1 << task)) continue;     // task already taken
            int newMask = mask | (1 << task);
            int newCost = dp[mask] + cost[worker][task];
            if (newCost < dp[newMask]) dp[newMask] = newCost;
        }
    }
    return dp[fullMask];
}

__builtin_popcount(mask) counting the set bits is what recovers "which worker are we assigning next" without needing a second dimension in the DP table — since exactly popcount(mask) workers must have been assigned to reach a state with popcount(mask) tasks used. That collapses this problem's state from TSP's O(2ⁿ · n) down to O(2ⁿ), and the transition cost from O(n) to O(n) per mask but without the extra last dimension — an O(2ⁿ · n) algorithm overall, one factor of n cheaper than TSP.

3. Digit DP: Counting Numbers with a Property in a Range

Digit DP answers questions like "how many numbers in [1, N] have digit sum divisible by 7?" — where N can be up to 10¹⁸, far too large to check one number at a time. The state is built from the digits of N instead of the number's value: process one digit position at a time, tracking whether the number built so far is still "tight" against N's corresponding prefix (constraining which digits are legal next) or already strictly smaller (in which case any digit is legal).

#include <string>
#include <vector>
using namespace std;

class DigitCounter {
public:
    // Count integers in [0, N] whose digit sum is divisible by `divisor`
    long long countWithDigitSumDivisible(long long N, int divisor) {
        digits = to_string(N);
        this->divisor = divisor;
        memo.assign(digits.size(), vector<vector<long long>>(
            divisor, vector<long long>(2, -1)));
        return solve(0, 0, true);
    }

private:
    string digits;
    int divisor;
    // memo[position][digitSumSoFar % divisor][isTight] -- -1 means "not computed yet"
    vector<vector<vector<long long>>> memo;

    long long solve(int pos, int sumMod, bool tight) {
        if (pos == (int)digits.size()) {
            return sumMod == 0 ? 1 : 0;   // reached the end: does the digit sum qualify?
        }
        if (!tight && memo[pos][sumMod][0] != -1) {
            return memo[pos][sumMod][0];
        }

        int limit = tight ? (digits[pos] - '0') : 9;
        long long count = 0;

        for (int d = 0; d <= limit; d++) {
            bool nextTight = tight && (d == limit);
            count += solve(pos + 1, (sumMod + d) % divisor, nextTight);
        }

        if (!tight) memo[pos][sumMod][0] = count;   // only cache the "free" (non-tight) states
        return count;
    }
};

Memoization only applies to tight == false states: once a number's prefix is strictly smaller than N's, the remaining choices don't depend on N at all, so that subproblem's answer is reusable across every branch that reaches the same (position, sumMod) with slack already used up. tight == true states are never reusable — they're only ever reached along the single path that exactly matches N's prefix, so there's at most one such state per digit position and caching them buys nothing. This is what keeps the whole search to O(digits · divisor) distinct memoized states, each doing O(10) work, despite N itself being astronomically large.

4. Recognizing State-Compression Problems

Bitmask DP and digit DP are both instances of the same underlying move: when a problem's naive state space is enormous but the actual number of distinct states that matter is small, find a compact encoding for that smaller space and index a DP table by it directly. Two signals point to bitmask DP specifically: the number of "items" involved is small (typically n ≤ 20, since 2ⁿ needs to stay computationally reasonable), and the problem cares about which subset has been used, not the order within it. Digit DP's signal is different: the input is a huge numeric range and the property being counted depends on the number's digits (digit sum, count of a specific digit, no two adjacent digits equal) rather than its magnitude directly.

5. Hands-on Exercise

6. Knowledge Check

Four quick questions. Expand each to check your answer.