Algorithms

"""backtracking, trying the N queens problem. more:
- https://en.wikipedia.org/wiki/Backtracking
- https://en.wikipedia.org/wiki/Eight_queens_puzzle

takeaways from backtracking:

1.  I always called the technique  "recursive backtracking", but most just
    call it "backtracking" because recursion is so commonly used. (could do
    iter + stack if wanted.)

2. backtracking feels morally similar to brute force. (still seems exponential
   or super-exponential.) though cutting off whole search subtrees is
   undoubtedly good. (e.g., if can't place 6th queen from these 5, no need to
   also try placing 7th and 8th)

3. even simple heuristics dramatically reduce search space. for example,
   knowing there's 1 queen per row, so searching row-by-row, enormously
   reduces the search space.

some things i learned with display: - clearing screen is really easy - seeing
the attempts is really fun - doing so makes attempts 2-3 ORDERS OF MAGNITUDE
slower
    - e.g., 300 vs 270,000 attempts per second
- as always, measure before optimizing is good
    - (this i/o bound was sufficient to remove for what i wanted to try)

better implementations would probably use bitmaps and vector ops to check
row/col/diag collisions, rather than iterating like I do here

one surprising thing is that stupider searching (trying each queen placement
from [0][0] through [n][n]) checks 4x attempts/second vs smarter search
(automatically assigning queen i to row i, and only searching [i][0] through
[i][n])
"""

import os
import time
from typing import Optional

attempt = 0

last_reported = time.time()
attempts_since_last_report = 0
report_interval_s = 1


def render(board_size: int, queens: set[tuple[int, int]]):
    global attempt
    attempt += 1
    for row in range(board_size):
        buf = []
        for col in range(board_size):
            if (row, col) in queens:
                buf.append("[Q]")
            else:
                buf.append("[ ]")
        print("".join(buf))
    print(f"{len(queens)} queens (attempt {attempt})")


def render_lite(board_size: int, queens: set[tuple[int, int]]):
    global attempt, last_reported, attempts_since_last_report
    attempt += 1
    attempts_since_last_report += 1
    now = time.time()
    delta = now - last_reported
    if delta > report_interval_s:
        last_reported = now
        clear()
        aps = int(attempts_since_last_report / delta)
        attempts_since_last_report = 0
        print(
            f"attempt {attempt} (@{len(queens)}/{board_size} queens) ({aps} attempts/s)"
        )


def clear():
    os.system("cls" if os.name == "nt" else "clear")
    # print()
    pass


def tick():
    # time.sleep(0.01)
    pass


def collides(board_size: int, queens: set[tuple[int, int]]):
    for queen in queens:
        in_row = len([q for q in queens if q[0] == queen[0]])
        if in_row > 1:
            # print(f"{queen} collides row {queen[0]}")
            return True
        in_col = len([q for q in queens if q[1] == queen[1]])
        if in_col > 1:
            # print(f"{queen} collides col {queen[1]}")
            return True
        # 2 diagonal checks.
        # (1) UL --> BR
        r, c = queen
        n_ulbr = 0
        while r >= 0 and c >= 0:
            if (r, c) in queens:
                n_ulbr += 1
            r -= 1
            c -= 1
        r, c = queen[0] + 1, queen[1] + 1
        while r < board_size and c < board_size:
            if (r, c) in queens:
                n_ulbr += 1
            r += 1
            c += 1
        if n_ulbr > 1:
            # print(f"{queen} collides UL -> BR")
            return True
        # (1) BL --> UR
        # - to BL: rows increase, columns decrease
        r, c = queen
        n_blur = 0
        while r < board_size and c >= 0:
            if (r, c) in queens:
                n_blur += 1
            r += 1
            c -= 1
        # - to UR: rows decrease, columns increase
        r, c = queen[0] - 1, queen[1] + 1
        while r >= 0 and c < board_size:
            if (r, c) in queens:
                # print(f"({r}, {c}) in queens")
                n_blur += 1
            r -= 1
            c += 1
        if n_blur > 1:
            # print(f"{queen} collides BL -> UR")
            return True
    return False


def test_collides():
    assert collides(8, set([(0, 0), (0, 1)]))
    assert collides(8, set([(0, 0), (1, 0)]))
    assert collides(8, set([(0, 0), (1, 1)]))

    assert collides(8, set([(7, 7), (7, 0)]))
    assert collides(8, set([(7, 7), (7, 6)]))
    assert collides(8, set([(7, 7), (0, 7)]))
    assert collides(8, set([(7, 7), (6, 7)]))
    assert collides(8, set([(7, 7), (0, 0)]))

    assert collides(8, set([(3, 3), (2, 2)]))
    assert collides(8, set([(3, 3), (4, 4)]))
    assert collides(8, set([(3, 3), (4, 2)]))
    assert collides(8, set([(3, 3), (2, 4)]))

    assert not collides(8, set([(0, 0), (1, 2)]))
    assert not collides(8, set([(0, 0), (2, 1)]))
    assert not collides(8, set([(0, 0), (2, 3)]))


def place_queen_naive(
    board_size: int, queens: set[tuple[int, int]]
) -> Optional[set[tuple[int, int]]]:
    """with rendering, this checked like ~300 attempts per second, which is hilariously
    low"""
    # print(f"place_queen(board_size={board_size}, queens={len(queens)})")
    # base case: we're done when we've placed board_size queens
    if len(queens) == board_size:
        return queens
    # brute force w/ .... backtracking?
    for r in range(board_size):
        for c in range(board_size):
            if (r, c) in queens:
                continue
            attempt = queens | {(r, c)}

            # ~300 aps
            # tick()
            # clear()
            # render(board_size, attempt)

            # ~270,000 aps
            render_lite(board_size, attempt)

            if not collides(board_size, attempt):
                # attempt is viable, so proceed
                continuation = place_queen_naive(board_size, attempt)
                if continuation is not None:
                    return continuation  # full solution
                # otherwise, we just continue (loops to next C)
    return None


def place_queen_byrow(
    board_size: int, queens: set[tuple[int, int]]
) -> Optional[set[tuple[int, int]]]:
    """with rendering, this checked like ~300 attempts per second, which is hilariously
    low"""
    # print(f"place_queen(board_size={board_size}, queens={len(queens)})")
    # base case: we're done when we've placed board_size queens
    if len(queens) == board_size:
        return queens
    # by row, still brute force w/ .... backtracking?
    r = len(queens)  # 1 queen / row (else trivial collision)
    for c in range(board_size):
        if (r, c) in queens:
            continue
        attempt = queens | {(r, c)}

        # ~300 aps
        # tick()
        # clear()
        # render(board_size, attempt)

        # ~70,000 aps
        render_lite(board_size, attempt)

        if not collides(board_size, attempt):
            # attempt is viable, so proceed
            continuation = place_queen_byrow(board_size, attempt)
            if continuation is not None:
                return continuation  # full solution
            # otherwise, we just continue (loops to next C)
    return None


def main() -> None:
    test_collides()

    # placements = place_queen_naive(8, set())  # solution @ 1,502,346
    # placements = place_queen_byrow(8, set())  # few seconds, solution @ 876
    placements = place_queen_byrow(20, set())  # ~1min, solution @ 3,992,511
    if placements is not None:
        render(len(placements), placements)
    print(f"Complete, result was: {placements}")

    # board_size = 8
    # clear()
    # render(board_size, set([(1, 2), (7, 7)]))
    # tick()
    # clear()
    # render(board_size, set([(0, 0), (0, 7)]))


if __name__ == "__main__":
    main()

From my source repo.

"""AKA cumulative sum, scan.

For an input array x, its prefix sum is a new array y, where each entry y[i]
is the cumulative sum of x from x[0] through x[i].
"""

prefix_sum_oneliner = lambda x: [sum(x[:i]) for i in range(len(x))]
"""cute O(n^2) prefix sum that's a python one-liner."""


def prefix_sum(xs: list[int]) -> list[int]:
    """Returns prefix sum of xs in O(n)."""
    y = [0] * len(xs)
    for i, x in enumerate(xs):
        y[i] = x + (y[i - 1] if i > 0 else 0)
    return y


def main() -> None:
    a = [1, 2, 3, 4, 5]
    b = prefix_sum(a)
    print(a)
    print(b)


if __name__ == "__main__":
    main()

From my source repo.

"""Kadane's algorithm for finding the maximum sum sub-array in O(n)."""

from typing import Optional


def kadane_no_indices(xs: list[int]) -> Optional[int]:
    """easy mode: just returns best sum (or None if empty)"""
    best_sum, cur_sum = None, 0
    for x in xs:
        cur_sum = max(cur_sum + x, x)  # continue or restart
        best_sum = max(best_sum, cur_sum) if best_sum is not None else cur_sum
    return best_sum


def kadane(xs: list[int]) -> tuple[int, int, Optional[int]]:
    """returns (best_start, best_end (inclusive!), best_sum | None if empty)"""
    cur_start, cur_end = 0, 0
    best_start, best_end = 0, 0
    cur_sum = 0
    best_sum = None
    for i, x in enumerate(xs):
        if x > cur_sum + x:
            # starting new
            cur_start = i
            cur_end = i
            cur_sum = x
        else:
            # continue sum
            cur_sum = cur_sum + x
            cur_end = i

        # if new best, also update saved indices
        if best_sum is None or cur_sum > best_sum:
            best_start = cur_start
            best_end = cur_end
            best_sum = cur_sum

    return best_start, best_end, best_sum


def main() -> None:
    a = [-2, 1, -3, 4, -1, 2, 1, -5, 4]
    best_start, best_end, best_sum = kadane(a)
    print(a)
    print(f"best: sum(a[{best_start}:{best_end}]) (inclusive)")
    print(f"    = sum({a[best_start : best_end + 1]})")
    print(f"    = {sum(a[best_start : best_end + 1])}")
    print(f"    = {best_sum}")
    print(f"    = {kadane_no_indices(a)}")


if __name__ == "__main__":
    main()

From my source repo.

"""Quicksort:
- choose pivot element
- partition into < and > the pivot
- put pivot right at that boundary
- quicksort each partition

Key properties:
- worst case: O(n^2)
- average case: O(n log n)
- in-place
- space: O(log n) b/c you gotta have a O(log n) call stack
    - wikipedia says you gotta use "Sedgewick's trick" to sort the larger
      partition using tail recursion or iteration, else you might use O(n)
      recursive calls. hmm.

Quicksort's degenerate case is on an already sorted array:
- say we choose the pivot to be the last element
- the "< pivot" boundary moves from 0 to the last element (O(n))
- then, the left partition has n-1 elements, and the right partition is empty
- we quicksort the left partition, which happens O(n) times recursively
"""

import random


def quicksort(x: list[int], start: int = 0, end: int = -1) -> None:
    """in-place quicksorts x from start through end, both inclusive."""
    if end == -1:
        end = len(x) - 1

    # base case: any sequence of 0 or 1 (or negative, lol) elements is sorted
    if start >= end:
        return

    # conceptual notes:
    # - boundary is inclusive
    # - i >= boundary, i.e., we'll iterate at least as fast as the boundary grows
    # - we compare up to pivot (exclusive)
    pivot = x[end]  # subject of comparison
    boundary = start - 1  # invariant: all els <= boundary will be <= pivot
    for i in range(start, end):
        if x[i] <= pivot:
            # if x[i] <= pivot, it should live in the left (<= boundary)
            # region. we can accomplish this by:
            # - extending the boundary by one to accommodate it, then
            # - swapping it with whatever is there. this is ok b/c:
            #      - conceptually, we're swapping left (i >= boundary)
            #      - so either now i == boundary (and swap is noop),
            #      - or i > boundary, and what we're swapping x[i] with at
            #        x[boundary] has already been checked, and is > pivot
            #        anyway, so it's fine staying right of the boundary, moving
            #        potentially much further right to wherever x[i] is
            boundary += 1
            x[i], x[boundary] = x[boundary], x[i]
        else:
            # x[i] > pivot, so it should live right of the boundary. this is
            # fine, we just keep boundary where it is.
            pass

    # now, we have two partitions, and the pivot:
    # x[start : boundary + 1] (i.e., boundary-inclusive) that's <= pivot
    # x[boundary + 1 : end] (i.e., pivot-exclusive) that's > pivot
    # x[end] the pivot
    #
    # if we swap the pivot into right after the boundary, it'll live at the right place.
    # (bounds ok? yes: boundary is at most the same as end.)
    x[end], x[boundary + 1] = x[boundary + 1], x[end]

    # now the partitions and the pivot are at:
    # x[start : boundary + 1]  <-- <= pivot
    # x[boundary + 1]          <-- pivot
    # x[boundary + 2 : end]    <-- > pivot
    #
    # we quicksort the partitions
    quicksort(x, start, boundary)  # not `boundary + 1`, b/c inclusive
    quicksort(x, boundary + 2, end)  # boundary + 2 ok? just check in call


def main() -> None:
    n = 10
    x = random.sample(range(1, n + 1), n)
    print(f"Before: {x}")
    quicksort(x)
    print(f"After:  {x}")
    assert x == sorted(x)


if __name__ == "__main__":
    main()

From my source repo.

"""Mergesort

- divide
- conquer (sort)
- merge (easy b/c sorted)

Properties:
- best O(n log n)
- worst O(n log n)
- in general, not in-place;
- in general, needs O(n) extra space
"""

import random


def mergesort(x: list[int]) -> list[int]:
    # 0- or 1-el lists are sorted
    if len(x) <= 1:
        return x

    # divide. as with all binary stuff, OBOB easy to do.
    # n    n // 2   x[:middle], x[middle:]
    # ---  ---      ---
    # 0     0       n/a! would have returned
    # 1     0       n/a! would have returned
    # 2     1       [0], [1]
    # 3     1       [0], [1,2]
    # 4     2       [0,1], [2,3]
    # 5     2       [0,1], [2,3,4]
    # ...
    middle = len(x) // 2

    # conquer
    left = mergesort(x[:middle])
    right = mergesort(x[middle:])

    # merge
    sorted = []
    i, j = 0, 0
    while i < len(left) or j < len(right):
        # if either list is exhausted, add rest from the other. (might as well
        # do all at once.)
        if i == len(left):
            sorted.extend(right[j:])
            j = len(right)
        elif j == len(right):
            sorted.extend(left[i:])
            i = len(left)
        else:
            # standard case, we add whichever is smaller
            if left[i] < right[j]:
                sorted.append(left[i])
                i += 1
            else:
                sorted.append(right[j])
                j += 1
    return sorted


def main() -> None:
    n = 10
    x = random.sample(range(1, n + 1), n)
    print(f"Before: {x}")
    sorted_x = mergesort(x)
    print(f"After:  {sorted_x}")
    assert sorted_x == sorted(x)


if __name__ == "__main__":
    main()

From my source repo.

r"""Minimal generic binary tree (node) used in both DFS and BFS.

TreeNode.build() builds an array from a complete binary tree
representation. `None`s are fine, just not as parents.
Using 0-based indexing.

           0
        /     \
      1          2
    /  \       /   \
   3    4     5     6
  / \  / \   / \   / \
 7  8 9 10  11 12 13 14

node i,'s children are:
    left: i*2 + 1
    right: i*2 + 2
"""

from typing import Generic, TypeVar, Optional

T = TypeVar("T")


class TreeNode(Generic[T]):
    __slots__ = ("val", "left", "right")

    def __init__(
        self,
        val: Optional[T] = None,
        left: Optional["TreeNode"] = None,
        right: Optional["TreeNode"] = None,
    ):
        self.val = val
        self.left = left
        self.right = right

    def __repr__(self) -> str:
        l = self.left.val if self.left is not None else "None"
        r = self.right.val if self.right is not None else "None"
        return f"TreeNode({self.val}, L={l}, R={r})"

    @staticmethod
    def build(complete: list[Optional[T]]) -> "TreeNode[T]":
        """Returns root or errors if not possible (just to minimize None checks)."""
        if len(complete) == 0 or complete[0] is None:
            raise ValueError(f"Must pass root. Complete was: {complete}")

        root = TreeNode(complete[0])
        flat_nodes: list[Optional["TreeNode[T]"]] = [None] * len(complete)
        for i, val in enumerate(complete):
            if i == 0:
                flat_nodes[0] = root
                continue
            if val is None:
                continue
            parent_idx = (i - 1) // 2
            parent = flat_nodes[parent_idx]
            if parent is None:
                raise ValueError(
                    f"Node {i} ({val})'s parent {parent_idx} must exist, but was None"
                )
            node = TreeNode(val=val)
            if (parent_idx * 2) + 1 == i:
                parent.left = node
            else:
                assert (parent_idx * 2) + 2 == i
                parent.right = node
            flat_nodes[i] = node
        return root


def main() -> None:
    root = TreeNode.build(list(range(15)))
    print("root:", root)


if __name__ == "__main__":
    main()

From my source repo.

r"""DFS: Depth-first search.

Using 0-based indexing.

           0
        /     \
      1          2
    /  \       /   \
   3    4     5     6
  / \  / \   / \   / \
 7  8 9 10  11 12 13 14

node i,'s children are:
    left: i*2 + 1
    right: i*2 + 2

DFS order traversal use cases:

- pre-order:  renders topologically sorted (for every edge (u,v), u comes before v).
              good for copying tree. also, prefix notation.

- in-order:   renders a BST in sorted order

- post-order: render postfix notation (e.g., if internal operators & leaf values).
              also good for deleting tree (removes children first).
"""

from typing import Literal, Optional

from binary_tree import TreeNode, T

Order = Literal["pre", "in", "post"]
ALL_ORDERS: list[Order] = ["pre", "in", "post"]


def dfs_recursive_v1(
    node: TreeNode[T],
    order: Order,
    results: Optional[list[Optional[T]]] = None,
) -> list[Optional[T]]:
    """DFS, depth-first, recursive, kind of gross v1 mutating implementation

    I passed `results` as mutable state into recursive calls. I think
    this is less elegant than using what's returned by the recursion
    as the result.
    """
    # We default argument None, not [], b/c otherwise the default gets mutated!
    if results is None:
        results = []

    if order == "pre":
        results.append(node.val)

    if node.left is not None:
        # NOTE: No need to assign or extend, because
        # - assign? no, we're passing `results` which gets mutated
        # - extend? no, each node's traversal adds itself, so no need to multi-add
        dfs_recursive_v1(node.left, order, results)

    if order == "in":
        results.append(node.val)

    if node.right is not None:
        dfs_recursive_v1(node.right, order, results)

    if order == "post":
        results.append(node.val)

    return results


def dfs(node: Optional[TreeNode[T]]) -> list[Optional[T]]:
    """DFS, depth-first, minimal pure recursive implementation (pre-order)."""
    if node is None:
        return []
    return [node.val, *dfs(node.left), *dfs(node.right)]


def dfs_recursive(
    node: Optional[TreeNode[T]], order: Order
) -> list[Optional[T]]:
    """DFS, depth-first, any-order recursive implementation."""
    if node is None:
        return []
    return [
        *([node.val] if order == "pre" else []),
        *dfs_recursive(node.left, order),
        *([node.val] if order == "in" else []),
        *dfs_recursive(node.right, order),
        *([node.val] if order == "post" else []),
    ]


def dfs_i(root: TreeNode[T]) -> list[Optional[T]]:
    """DFS, depth-first, minimal iterative implementation (pre-order)"""
    res: list[Optional[T]] = []
    stack: list[Optional[TreeNode[T]]] = [root]
    while len(stack) > 0:
        node = stack.pop()
        if node is None:
            continue
        res.append(node.val)
        stack.append(node.right)
        stack.append(node.left)

    return res


def dfs_iterative(root: TreeNode[T], order: Order) -> list[Optional[T]]:
    """DFS, depth-first, iterative

    we add values back to the stack to control the traversal order. this two-pass
    approach lets us decide explicitly when to extract a node's value.
    """
    results = []
    stack: list[TreeNode[T] | Optional[T]] = [root]
    while len(stack) > 0:
        node_or_value = stack.pop()

        # we allow None to simplify adding vals (Optional[T])
        if node_or_value is None:
            continue

        # if we got a value on the stack, that means we should add it
        # to the results.
        if not isinstance(node_or_value, TreeNode):
            val = node_or_value
            results.append(val)
            continue

        # remember to add things "backwards" because stack will pop latest first
        node = node_or_value
        if order == "post":
            stack.append(node.val)
        if node.right is not None:
            stack.append(node.right)
        if order == "in":
            stack.append(node.val)
        if node.left is not None:
            stack.append(node.left)
        if order == "pre":
            stack.append(node.val)
    return results


def main() -> None:
    root = TreeNode.build(list(range(15)))

    # DFS - depth-first
    for order in ALL_ORDERS:
        recursive_v1 = dfs_recursive_v1(root, order)
        recursive = dfs_recursive(root, order)
        iterative = dfs_iterative(root, order)
        assert recursive_v1 == recursive
        assert recursive == iterative
        print(f"DFS ({order}):", recursive)

    # also check the minimal implementations are correct
    assert dfs(root) == dfs_recursive(root, "pre")
    assert dfs_i(root) == dfs_recursive(root, "pre")


if __name__ == "__main__":
    main()

From my source repo.

r"""BFS: Breadth-first search.

Using 0-based indexing.

           0
        /     \
      1          2
    /  \       /   \
   3    4     5     6
  / \  / \   / \   / \
 7  8 9 10  11 12 13 14

node i,'s children are:
    left: i*2 + 1
    right: i*2 + 2

BFS:
- traverses in complete-tree-as-array order (top to bottom by level, and left to right
  within each level)

BFS, interestingly:
- no recursion: doesn't have a natural recursive implementation
- no order: doesn't really have different order traversals (could do right before left,
            which reverses the within-level order)
"""

from collections import deque
from typing import Optional

from binary_tree import TreeNode, T


def bfs(
    root: TreeNode[T],
    results: Optional[list[Optional[T]]] = None,
) -> list[Optional[T]]:
    """BFS, breadth-first, iterative"""
    results = []
    queue: deque[Optional[TreeNode]] = deque([root])
    while len(queue) > 0:
        node = queue.popleft()
        if node is None:
            continue
        results.append(node.val)
        queue.append(node.left)
        queue.append(node.right)

    return results


def main() -> None:
    root = TreeNode.build(list(range(15)))

    # BFS - breadth-first
    print("BFS:", bfs(root))


if __name__ == "__main__":
    main()

From my source repo.