Backtracking with pruning

Why prune?

Plain backtracking builds every path to the end. Often you can tell much earlier that a partial path is hopeless: the running total already passed the target, or two queens already attack each other. Pruning is the act of returning immediately from such a call. The answers do not change; you just stop exploring branches that cannot produce any. The worst case Big-O may look the same, but in practice the number of calls can drop from millions to thousands.

Combination sum: reuse allowed

Given distinct positive candidates and a target, list every combination that adds up to the target, where each candidate can be used any number of times. To avoid listing [2, 3, 2] and [2, 2, 3] separately, build each combination in non-decreasing order: pass the current index down, not i + 1, because the same candidate may be used again. Sort the candidates so you can prune: once a candidate is bigger than what remains, every later one is too, so break out of the loop.

def combination_sum(candidates, target):
    candidates = sorted(candidates)
    result, path = [], []

    def backtrack(start, remaining):
        if remaining == 0:
            result.append(path[:])
            return
        for i in range(start, len(candidates)):
            c = candidates[i]
            if c > remaining:          # prune: this and all later ones are too big
                break
            path.append(c)
            backtrack(i, remaining - c)   # i, not i + 1: c may be reused
            path.pop()

    backtrack(0, target)
    return result

print(combination_sum([2, 3, 6, 7], 7))   # Output: [[2, 2, 3], [7]]

Letter combinations of a phone number

On an old phone keypad, 2 is "abc", 3 is "def", 4 "ghi", 5 "jkl", 6 "mno", 7 "pqrs", 8 "tuv" and 9 "wxyz". Given digits such as "23", list every word the keys could type: ad, ae, af, bd, … , cf. Each level of the tree is one digit and its branches are that digit's letters, so there are up to 4ⁿ results for n digits. There is nothing to prune, but it is the same choose-explore-un-choose shape. Watch the edge case: empty input gives an empty list, not [""].

Generating balanced parentheses

List every string of n pairs of brackets that is balanced, like "(())" and "()()" for n = 2. Instead of generating all 2²ⁿ strings and checking each, prune while building: you may add "(" only while fewer than n have been used, and ")" only while there are more open brackets than closed ones. Every completed string is then valid by construction. The number of answers is the nth Catalan number: 1, 1, 2, 5, 14, 42, … for n = 0, 1, 2, 3, 4, 5.

def generate_parentheses(n):
    result = []

    def backtrack(path, opened, closed):
        if len(path) == 2 * n:
            result.append("".join(path))
            return
        if opened < n:
            path.append("("); backtrack(path, opened + 1, closed); path.pop()
        if closed < opened:
            path.append(")"); backtrack(path, opened, closed + 1); path.pop()

    backtrack([], 0, 0)
    return result

print(generate_parentheses(3))
# Output: ['((()))', '(()())', '(())()', '()(())', '()()()']

N-Queens

Place n chess queens on an n × n board so that no two share a row, a column or a diagonal. Put exactly one queen in each row, going down row by row; at each row, try each column. Pruning makes it fast: keep three sets: used columns, used "r − c" diagonals (cells on one down-right diagonal share r − c) and used "r + c" anti-diagonals. A column is safe only if all three checks pass, so most branches end within a few rows. For n = 8 there are 92 solutions; for n = 2 and n = 3 there are none.

def n_queens_count(n):
    cols, diag, anti = set(), set(), set()

    def place(r):
        if r == n:
            return 1
        total = 0
        for c in range(n):
            if c in cols or (r - c) in diag or (r + c) in anti:
                continue                   # prune: attacked
            cols.add(c); diag.add(r - c); anti.add(r + c)
            total += place(r + 1)
            cols.remove(c); diag.remove(r - c); anti.remove(r + c)
        return total

    return place(0)

print(n_queens_count(4), n_queens_count(8))   # Output: 2 92

Sudoku: the classic example

A Sudoku solver uses the same plan. Find an empty cell, try digits 1 to 9, and keep a digit only if it does not already appear in that cell's row, column or 3 × 3 box (sets make each check O(1)). Recurse to the next empty cell; if no digit fits somewhere later, undo and try the next digit. Good solvers prune harder by filling first the cell with the fewest possible digits. You do not need to code it here, but you should be able to describe it in an interview.

Common mistakes

• Pruning that drops valid answers, for example break when the candidates are not sorted. Fix: sort first, or use continue instead of break.
• Using i + 1 in combination sum when reuse is allowed (or i when it is not). Fix: decide whether an item can repeat, and pass i or i + 1 to match.
• Forgetting to remove a queen's column and diagonals after the recursive call. Fix: undo exactly what you added, in the same call.
• Returning [""] for empty digits in letter combinations. Fix: handle the empty input first and return [].
• Checking validity only at the end, which wastes time on dead branches. Fix: check as early as possible, before recursing.