Shows all boxes that can be filled in a Picross puzzle given the input pattern(s).

picross_line_solver.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from argparse import ArgumentParser, RawDescriptionHelpFormatter


def solve(size: int, line: list[int]) -> list[bool]:
    solution = [False] * size

    needed_to_fill = sum(line) + len(line) - 1
    if needed_to_fill > size:
        return solution

    empty = size - needed_to_fill
    if empty > max(line):
        return solution

    current_pos = 0
    for number in line:
        current_pos += number
        intersect = number - empty
        while intersect > 0:
            solution[current_pos - intersect] = True
            intersect -= 1
        current_pos += 1

    return solution


def pretty_print_line(
    lines: list[list[int]], solutions: list[list[bool]], transpose: bool
) -> str:
    output = []
    for line, sol in zip(lines, solutions):
        output.append(
            " ".join(map(str, line))
            + " "
            + " ".join("■" if i else "□" for i in sol)
        )

    if transpose:
        empty = (" ",) * len(lines)
        print("\n".join(" ".join(i) for i in zip(*output) if i != empty))
    else:
        print("\n".join(output))


if __name__ == "__main__":
    PROG_NAME = "picross_line_solver"
    DESCRIPTION = [
        "Shows all boxes that can be filled in a Picross puzzle given the",
        "input pattern(s).",
        "",
        "$ {} -s 15 -i 2 7 1 -i 6 2 1 -i 4 4 3".format(PROG_NAME),
        "2 7 1 □ □ □ □ □ □ ■ ■ ■ ■ □ □ □ □ □",
        "6 2 1 □ □ □ □ ■ ■ □ □ □ □ □ □ □ □ □",
        "4 4 3 □ □ ■ ■ □ □ □ ■ ■ □ □ □ ■ □ □",
    ]

    PARSER = ArgumentParser(
        prog=PROG_NAME,
        formatter_class=RawDescriptionHelpFormatter,
        description="\n".join(DESCRIPTION),
    )
    PARSER.add_argument(
        "-s",
        "--size",
        help="length of the nonogram line",
        type=int,
        required=True,
    )
    PARSER.add_argument(
        "-i",
        "--input",
        nargs="+",
        action="append",
        help="layout of the line (i.e. a list of numbers)",
        type=int,
        required=True,
    )
    PARSER.add_argument(
        "-c",
        "--column",
        action="store_true",
        help="print the output vertically",
    )

    ARGS = PARSER.parse_args()
    SOLUTIONS = [solve(ARGS.size, line) for line in ARGS.input]
    pretty_print_line(ARGS.input, SOLUTIONS, ARGS.column)