Tree, digit and bitmask DP

Three shapes come up often enough to be worth recognising on sight, and each is announced by its constraints rather than by its wording. A rooted structure means tree DP. A bound like 10^18 with "how many numbers" means digit DP. An n of 20 means a bitmask. These are recognition problems first and implementation problems second.

The constraint says	The shape
the input is a tree, answer defined per subtree	tree DP
count values in [1, N] with N up to 10^18	digit DP
n ≤ 20, and the answer needs a set or an order	bitmask DP

Tree, digit and bitmask DP side by side: a child returning a pair up to its parent, the digit-DP state tuple branching on the tight flag, and a subset held as bits of an integer100%−+FitSVGEsc

Scroll to zoom · drag to pan · 0 fits · Esc closes

Tree DP: postorder, returning a tuple

The states are the nodes and the dependency is the tree itself: a node needs its children's answers, so the fill order is postorder. One traversal visits each node once, so the whole thing is O(n).

What each call returns is the design decision. Almost always it is a small tuple — one entry per way the node can participate. For the maximum sum of nodes with no two adjacent chosen, that is two numbers:

def solve(node):
    """(best with this node taken, best with it skipped)"""
    if node is None:
        return (0, 0)
    lt, ls = solve(node.left)
    rt, rs = solve(node.right)
    take = node.val + ls + rs          # children must be skipped
    skip = max(lt, ls) + max(rt, rs)   # children free to do either
    return (take, skip)

answer = max(solve(root))

skip takes the best of each child independently — a skipped parent does not force its children to be taken. The pattern generalises: one tuple entry per way a node can participate, combined at the parent by whatever links a node to its children. Diameter returns (height, best path through here); colouring returns one entry per colour.

Depth is the practical limit. A path-shaped tree of 10^5 nodes recurses 10^5 deep and hits Python's ~1,000-frame ceiling, so a real submission uses an explicit stack or an iterative postorder over a precomputed order.

Digit DP: counting with a tight flag

The question is "how many numbers from 1 to N have property P", with N far too large to enumerate — 10^18 values is beyond counting, but N has only 19 digits. So build the number one digit at a time, left to right, and count the completions.

The state is (position, tight, started, extra):

• position — how many digits are fixed so far.
• tight — whether the prefix so far is exactly N's prefix. This is the one people miss. While tight, the next digit may not exceed N's digit at that position, or the number overruns N; place anything smaller and every remaining position is free to be 0 through 9.
• started — whether a non-zero digit has appeared, which separates a leading zero from a genuine digit 0.
• extra — whatever the property needs to carry: a running remainder, the previous digit, a count.

Counting numbers up to N whose digit sum is divisible by 3:

from functools import lru_cache

def count(N):
    d = list(map(int, str(N)))

    @lru_cache(maxsize=None)
    def go(pos, tight, started, rem):
        if pos == len(d):
            return 1 if started and rem == 0 else 0
        top = d[pos] if tight else 9
        total = 0
        for digit in range(top + 1):
            total += go(pos + 1,
                        tight and digit == top,
                        started or digit > 0,
                        (rem + digit) % 3)
        return total

    return go(0, True, False, 0)

Size it with the usual product from what DP actually is: 19 positions x 2 tight x 2 started x 3 remainders = 228 states, each looping over at most 10 digits — about 2,300 operations to answer a question about 10^18 numbers. Ranges come free: the count over [L, R] is count(R) - count(L - 1).

The bug to expect is a memo keyed without tight: a count worked out with all ten digits available is not valid on a branch where the next digit is capped, and without tight in the key the two share an entry. Memoising only when tight is false sidesteps it, and those entries are the overwhelming majority anyway.

Bitmask DP: the subset is an integer

Bit i of an integer says whether item i is in the set, so a subset is an array index. mask & (1 << i) tests membership, mask | (1 << i) adds, and mask.bit_count() gives the size.

The ceiling is arithmetic, not taste. 2^20 is 1,048,576 states, which fits comfortably; 2^25 is 33 million, which does not once each state carries a transition loop. An n of about 20 in the constraints is a signal, not a coincidence, exactly as complexity by counting reads it.

Travelling salesman is the canonical version: dp[mask][last] is the cheapest route that has visited mask and is standing at last.

def tsp(dist):
    n = len(dist)
    INF = float('inf')
    dp = [[INF] * n for _ in range(1 << n)]
    dp[1][0] = 0                                    # start at city 0
    for mask in range(1 << n):
        for last in range(n):
            if dp[mask][last] == INF:
                continue
            for nxt in range(n):
                if mask & (1 << nxt):
                    continue
                m2 = mask | (1 << nxt)
                dp[m2][nxt] = min(dp[m2][nxt], dp[mask][last] + dist[last][nxt])
    return min(dp[(1 << n) - 1][k] + dist[k][0] for k in range(n))

That is 2^n x n states with n transitions each: at n = 20, 2^20 x 20 x 20 = 4 x 10^8 — fine compiled, far too slow in Python, where you want n nearer 15. Memory is 2^n x n cells, 21 million at n = 20, about 160 MB at 8 bytes each, which is why these problems cap n where they do.

When the transition splits a set into two parts rather than adding one element, you enumerate submasks:

sub = mask
while sub:
    # partition into sub and mask ^ sub
    sub = (sub - 1) & mask

The total over every mask is 3^n, not 4^n, because each element is in sub, in the complement, or outside mask altogether. At n = 16 that is 43 million, which runs; at n = 20 it is 3.5 x 10^9, which does not. Knowing which of 2^n and 3^n your loop costs is the difference between a solution and a timeout — the same size-first habit that decides the table in the knapsack family.

In an interview

Read the constraint aloud and name the shape from it: "n is 18 and the answer is an assignment of people to tasks — that is 2^n states over subsets." Doing that in the first minute is worth more than the implementation, because it is the step that cannot be recovered later.

For a tree, say what the tuple means before writing the traversal. For digit DP, name the four state components and point at tight as the one that makes the count exact.

The mistake that loses points: writing a bitmask solution without doing the arithmetic. 2^n at n = 30 is a billion states, and no amount of clean code makes that run. Say the number before you commit to the approach.

Check yourself

A problem gives n ≤ 22 cities and asks for the cheapest tour. How many states, and how many operations, and what does that tell you about the language?

2^22 x 22 ≈ 92 million states, each with up to 22 transitions — about 2 x 10^9 operations, and 92 million cells of memory. That is a compiled-language answer at best; in Python you would need the constraint to be nearer n = 15.

Your digit DP memoises on (pos, started, rem) and drops tight. count(999) returns 333, which is right, but count(527) returns 199 where the answer is 175. Name the missing component, say why 999 escapes, and decide whether refusing to cache the tight branch costs you anything.

tight is missing. The digit loop runs upward, so the free branch — digit below the cap — fills a key first, and the tight branch then hits that entry and inherits a count that allowed all ten digits, so the answer comes out too high. You can check the 175 by hand: a digit sum divisible by 3 means the number is, and there are 175 multiples of 3 up to 527. N = 999 escapes because every digit of N is 9, so capped and free allow the same digits. Refusing to cache while tight costs almost nothing: exactly one prefix is tight at each position, so that is 19 of the 228 states at N near 10^18, and at most 190 extra digit steps.

A node's subtree answer has to satisfy "no two chosen nodes adjacent". Decide what the postorder call returns, and say how many entries it needs for a k-colouring variant and what that costs.

Return a pair, (best with this node taken, best with it skipped), because the parent has to know how the child's answer used the child itself; one number cannot separate "took the child" from "skipped it", so the parent cannot enforce adjacency. The rule is one entry per way a node can participate, so k-colouring returns k entries, one per colour. Combining them is k choices at the parent times k at each child — O(n x k^2) — which drops to O(n x k) if each child also reports its second-best colour, since the parent only ever needs the best colour that is not its own.