Palindrome partition by backtracking in Python

Recursively tries each palindromic prefix and enumerates all partitions.

File: palindrome_partition_backtracking.py, 65 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

palindrome_partition_backtracking.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/palindrome-partition-backtracking/. The authors are listed in prov/8C52F0.jsonl.

"""All partitions of a string into palindromes, by backtracking with an explicit stack.

Indices are 0-based and s[i..j] includes both ends. Pieces may be odd or even
palindromes. Symbols are only tested for equality.
"""

from collections.abc import Iterator


def palindrome_ends(s: str) -> list[list[int]]:
    """Return ends with ends[i] the sorted list of j such that s[i..j] is a palindrome."""
    n = len(s)
    pal = [[False] * n for _ in range(n)]
    ends: list[list[int]] = [[] for _ in range(n)]
    for i in range(n - 1, -1, -1):  # @line L1
        for j in range(i, n):  # @line L2
            pal[i][j] = s[i] == s[j] and (j - i < 2 or pal[i + 1][j - 1])  # @line L3
            if pal[i][j]:  # @line L4
                ends[i].append(j)
    return ends


def palindrome_partitions(s: str) -> Iterator[list[str]]:
    """Yield every partition of s into nonempty palindromes, each one time.

    The partitions come in lexicographic order of their piece lengths, so the
    first one is the split into single symbols. The empty string has one
    partition, the empty list.
    """
    n = len(s)
    ends = palindrome_ends(s)
    path: list[str] = []
    stack: list[list[int]] = [[0, 0]]  # @line L5
    while stack:  # @line L6
        frame = stack[-1]  # @line L7
        i, k = frame
        if i == n:  # @line L8
            yield list(path)
            stack.pop()
            if i > 0:
                path.pop()
            continue
        if k == len(ends[i]):  # @line L9
            stack.pop()
            if i > 0:
                path.pop()
            continue
        frame[1] = k + 1  # @line L10
        j = ends[i][k]
        path.append(s[i : j + 1])  # @line L11
        stack.append([j + 1, 0])


def count_palindrome_partitions(s: str) -> int:
    """Return the number of palindromic partitions of s, by the prefix recurrence."""
    n = len(s)
    ends = palindrome_ends(s)
    ways = [0] * (n + 1)
    ways[n] = 1
    for i in range(n - 1, -1, -1):
        ways[i] = sum(ways[j + 1] for j in ends[i])
    return ways[0]

test_palindrome_partition_backtracking.py 67 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/palindrome-partition-backtracking/. The authors are listed in prov/8C52F0.jsonl.

import itertools
import random

from palindrome_partition_backtracking import (
    count_palindrome_partitions,
    palindrome_ends,
    palindrome_partitions,
)


def brute_partitions(s: str) -> list[list[str]]:
    """All partitions from all 2^(n-1) sets of cuts, kept if every piece is a palindrome."""
    n = len(s)
    if n == 0:
        return [[]]
    out = []
    for mask in range(1 << (n - 1)):
        cuts = [0] + [i + 1 for i in range(n - 1) if mask >> i & 1] + [n]
        pieces = [s[a:b] for a, b in itertools.pairwise(cuts)]
        if all(p == p[::-1] for p in pieces):
            out.append(pieces)
    return out


def test_example() -> None:
    assert list(palindrome_partitions("abaab")) == [
        ["a", "b", "a", "a", "b"],
        ["a", "b", "aa", "b"],
        ["a", "baab"],
        ["aba", "a", "b"],
    ]
    assert palindrome_ends("abaab") == [[0, 2], [1, 4], [2, 3], [3], [4]]


def test_edges() -> None:
    assert list(palindrome_partitions("")) == [[]]
    assert list(palindrome_partitions("a")) == [["a"]]
    assert list(palindrome_partitions("abcabc")) == [list("abcabc")]
    for n in range(1, 12):
        parts = list(palindrome_partitions("a" * n))
        assert len(parts) == 2 ** (n - 1)
        assert count_palindrome_partitions("a" * n) == 2 ** (n - 1)
    assert count_palindrome_partitions("") == 1


def test_random_against_brute_force() -> None:
    rng = random.Random(8052)
    for trial in range(900):
        k = (1, 2, 3)[trial % 3]
        n = rng.randrange(0, 12)
        s = "".join(rng.choice("abc"[:k]) for _ in range(n))
        got = list(palindrome_partitions(s))
        want = brute_partitions(s)
        assert len(got) == len(want)
        assert sorted(got) == sorted(want), s
        assert count_palindrome_partitions(s) == len(want)
        for p in got:
            assert "".join(p) == s


def test_long_input_no_deep_recursion() -> None:
    s = "ab" * 600
    first = next(palindrome_partitions(s))
    assert len(first) == len(s)

What the next level needs

the tests did not run in this build

Read the full page about Palindrome partition by backtracking