pyramid_puzzle.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional


@dataclass
class Node:
    d: int
    x: int
    value: Optional[int] = None
    parent_l: Optional[Node] = None
    parent_r: Optional[Node] = None

    def __str__(self) -> str:
        if self.value is not None:
            return str(self.value).rjust(2, " ")
        else:
            return "__"


class Tree:
    def __init__(self, depth: int, values_to_fill: list[int], start_nodes: dict[tuple[int, int], int]) -> None:
        self.depth = depth
        self.values_to_fill = values_to_fill
        self.to_solve: set[tuple[int, int]] = set()
        self.nodes = self.create_tree()
        for (d, x), v in start_nodes.items():
            self.assign_node(d, x, v)
        self.solve()

    def __str__(self) -> str:
        output = ""
        for d in range(self.depth):
            line = " ".join([str(self.nodes[(d, x)]) for x in range(d + 1)])
            output += line.center((self.depth) * 3) + "\n"
        return output

    def assign_node(self, d: int, x: int, v: int) -> None:
        self.to_solve.discard((d, x))
        self.nodes[(d, x)].value = v
        self.values_to_fill.remove(v)
        if d < (self.depth - 1):
            if self.nodes[(d + 1, x)].value is None:
                self.to_solve.add((d + 1, x))
            if self.nodes[(d + 1, x + 1)].value is None:
                self.to_solve.add((d + 1, x + 1))

    def create_tree(self) -> dict[tuple[int, int], Node]:
        tree = {}
        for d in range(self.depth):
            for x in range(d + 1):
                tree[(d, x)] = Node(d, x)
                if d == 0:
                    continue
                if x > 0:
                    tree[(d, x)].parent_l = tree[(d - 1, x - 1)]
                if x < d:
                    tree[(d, x)].parent_r = tree[(d - 1, x)]
        return tree

    def solve(self) -> None:
        while len(self.to_solve):
            for d, x in self.to_solve:
                solution = self.check_neighbor_plus_parent(self.nodes[(d, x)])
                if solution is not None:
                    self.assign_node(d, x, solution)
                    break
                else:
                    options = self.get_options(self.nodes[(d, x)])
                    if len(options) == 1:
                        self.assign_node(d, x, options.pop())
                        break

    def get_options(self, node) -> set[int]:
        def options(parent_v):
            return {v for v in self.values_to_fill if parent_v - v in self.values_to_fill}

        if node.parent_l is None or node.parent_l.value is None:
            return options(node.parent_r.value)
        elif node.parent_r is None or node.parent_r.value is None:
            return options(node.parent_l.value)
        else:
            return options(node.parent_l.value).intersection(options(node.parent_r.value))

    def check_neighbor_plus_parent(self, node) -> Optional[int]:
        if node.parent_l is not None and node.parent_l.value is not None:
            neighbor_l = self.nodes[(node.d, node.x - 1)]
            if neighbor_l.value is not None:
                return node.parent_l.value - neighbor_l.value
        if node.parent_r is not None and node.parent_r.value is not None:
            neighbor_r = self.nodes[(node.d, node.x + 1)]
            if neighbor_r.value is not None:
                return node.parent_r.value - neighbor_r.value
        return None


if __name__ == "__main__":
    values_to_use = [64, 18, 46, 37, 29, 21, 13, 9, 9, 8, 8, 8, 8, 8, 7, 6, 5, 5, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]
    tree = Tree(depth=7, values_to_fill=values_to_use, start_nodes={(0, 0): 64, (1, 0): 18})
    print(tree)