find-solutions.py

import sys
from operator import mul
from functools import reduce
from collections import Counter
import argparse
import csv


# What to put between operands when printing a multiplication or addition.
MULT = " . "
ADD = " + "


class Number:
    def __init__(self, n: int | None = None, counter: Counter | None = None):
        if n is None:
            assert counter is not None
            self.counter = counter
            self.value = reduce(
                mul, [prime**power for prime, power in counter.items()], 1
            )
        else:
            assert counter is None
            self.value = n
            self.counter = Counter(factors(n))

    def __str__(self) -> None:
        parts = []
        for prime, power in sorted(self.counter.items()):
            parts.append(f"{prime}" if power == 1 else f"{prime}^{power}")

        return f"{self.value} = {self.parts()}"

    def parts(self) -> str:
        parts = []
        for prime, power in sorted(self.counter.items()):
            parts.append(f"{prime}" if power == 1 else f"{prime}^{power}")

        return MULT.join(parts)

    def distribute(self, other: "Number") -> tuple["Number", "Number", "Number"]:
        """
        Given another number ('other') return x, y, and z where:

            self + other = x * (y + z)

        do this by pulling common prime factors (with the highest common power) out
        of self and other, and using that to form x. Then y and z have the factors
        in x removed.
        """
        common = Counter()
        for prime, power in sorted(self.counter.items()):
            if prime in other.counter:
                common[prime] = min(other.counter[prime], power)

        rest1 = Counter(self.counter) - common
        rest2 = Counter(other.counter) - common

        c, r1, r2 = Number(counter=common), Number(counter=rest1), Number(counter=rest2)

        # Sanity check we've done things correctly.
        assert self.value + other.value == c.value * (r1.value + r2.value)

        return c, r1, r2


def factors(n: int) -> list[int]:
    assert n > 0

    if n == 1:
        return [1]

    result = []

    while n % 2 == 0:
        result.append(2)
        n //= 2

    d = 3
    while d * d <= n:
        while n % d == 0:
            result.append(d)
            n //= d
        d += 2

    if n > 1:
        result.append(n)

    return result


def sum_of_squares(a: int, b: int) -> str:
    aa = Number(a * a)
    bb = Number(b * b)
    common, rest1, rest2 = aa.distribute(bb)
    return (
        f"{common.value} ({common.parts()}){MULT}{rest1.value}{ADD}{rest2.value} "
        f"({rest1.value} ({rest1.parts()}){ADD}{rest2.value} ({rest2.parts()}))"
    ).replace("1 ()", "1")


def p(s: str, a: int, b: int | None = None) -> None:
    if b is None:
        print(f"{s:4s}: {a:>7d} =", MULT.join(map(str, factors(a))))
    else:
        ss = a * a + b * b
        print(f"{s:4s}: {ss:>7d} =", sum_of_squares(a, b))


def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Print integers a and b that satisfy the puzzle condition. "
            "See https://www.youtube.com/watch?v=UoME4zN5NlI for details."
        ),
    )
    parser.add_argument(
        "--n", type=int, default=2000, help="The upper limit on a and b."
    )
    parser.add_argument("--tsv", action="store_true", help="Write output as TSV")

    return parser.parse_args()


def main():
    args = get_args()
    tsv = args.tsv
    if tsv:
        writerow = csv.writer(sys.stdout, delimiter="\t").writerow
        writerow(("a", "b", "ss", "ab+1", "quotient"))
    for a in range(1, args.n):
        for b in range(a, args.n):
            ss = a * a + b * b
            ab = a * b + 1
            if ss % ab == 0:
                q = ss // ab
                if tsv:
                    writerow((a, b, ss, ab, q))
                else:
                    p("a", a)
                    p("b", b)
                    p("ss", a, b)
                    p("ab+1", ab)
                    p("q", q)
                    print()


if __name__ == "__main__":
    main()