Linear equations that start with a given sequence of moves in the collatz conjecture

collatz.py

import typing
import math

def collatz(n: int) -> typing.List[int]:
    """Returns the Collatz sequence starting at n."""
    seq = [n]
    while n != 1:
        if n % 2 == 0:
            n = n // 2
        else:
            n = (3 * n + 1) // 2
        seq.append(n)
    return seq

class LinearEq:
    """
    ak + b
    """
    def __init__(self, a: int, b: int) -> None:
        self.a = a
        self.b = b

    def evaluate(self, k: int) -> int:
        return self.a * k + self.b

    def __str__(self) -> str:
        if self.b == 0:
            return f"{self.a}k"
        return f"{self.a}k + {self.b}"

    def scale(self, k: int) -> "LinearEq":
        return LinearEq(self.a * k, self.b * k)

    def translate(self, k: int) -> "LinearEq":
        return LinearEq(self.a, self.b + k)

def pattern_of_linear(eq: LinearEq) -> str:
    """
    Returns the pattern of the Collatz sequence produced by the given linear equation.
    """
    # log2 of eq.a
    exp = int(math.log2(eq.a))

    # todo: limit to `exp` iterations instead of chopping at the end
    seq = collatz(eq.evaluate(1))[:exp]

    return ''.join(['+' if n % 2 == 1 else '-' for n in seq])

def linear_of_pattern(pattern: str) -> LinearEq:
    """
    Returns the linear equation that produces the given pattern.
    """
    eq = LinearEq(1, 0)
    for c in pattern:
        # try both, not sure which one I need, can I figure that out ahead of time?
        #
        # k -> 2k
        even = LinearEq(2, 0).scale(eq.a).translate(eq.b)
        # k -> 2k+1
        odd = LinearEq(2, 1).scale(eq.a).translate(eq.b)

        if pattern.startswith(pattern_of_linear(even)):
            eq = even
        else:
            eq = odd

    return eq

eq = linear_of_pattern("++-+----+---+++++--+-++-+++")
print(eq)
for k in range(1, 10):
    start = eq.evaluate(k)
    seq = collatz(start)
    print(f"{k}\t{start}\t{''.join(['+' if n % 2 == 1 else '-' for n in seq])}")

# $ python collatz.py
# 134217728k + 76130123
# 1	210347851	++-+----+---+++++--+-++-++++---+-+--+-++---++----+-++++------+++-++---+++++--++-+++-+-++-++-+---+--++++-+++++---+-+-+---+--+++----+---+
# 2	344565579	++-+----+---+++++--+-++-+++----+++---+-+--+-----+---+-+-+++-+-++++-+--+-++---+++---+-+--+-+++++-+-++-+++-++++-+--+++-++-++++++--++++---+-+-+---+--+++----+---+
# 3	478783307	++-+----+---+++++--+-++-+++++-+++---++-++----+-++-+--+++++-+-+-++-+--+-+----++++-++++++-+++-++++-----+++--++++++--+++++--+--++-+-+--++-+++--+++-+--++-+--++-+++--++-+---++++--++---++--+++++---+--+---+--+-+++++-+-++-+++-++++-+--+++-++-++++++--++++---+-+-+---+--+++----+---+
# 4	613001035	++-+----+---+++++--+-++-+++-++----+-++--+-----++--+-++-+-+----+-+-++++-+--+-++---+++---+-+--+-+++++-+-++-+++-++++-+--+++-++-++++++--++++---+-+-+---+--+++----+---+
# 5	747218763	++-+----+---+++++--+-++-++++-+---+-+++--+-+-++-+++---+------------+++--++-+-----+++++-+-++-+++-++++-+--+++-++-++++++--++++---+-+-+---+--+++----+---+
# 6	881436491	++-+----+---+++++--+-++-+++--+++-+-+-++--+---+---+---++------++-----+--+--+--++-+--+---+
# 7	1015654219	++-+----+---+++++--+-++-++++++-+-----++--++--+--++-+++--+----+-+-+---+--++---+-+----+++-+-++-+++-++++-+--+++-++-++++++--++++---+-+-+---+--+++----+---+
# 8	1149871947	++-+----+---+++++--+-++-+++-+-++-++-+-+-+-++----++---+++--++-+--++-++++--++++--+-+-+-----+-+--+++----+-+--++--++-++++-+--+++-++-++++++--++++---+-+-+---+--+++----+---+
# 9	1284089675	++-+----+---+++++--+-++-++++--++-----++++--+-+--+--+++++------+-+----++++++-++--+----++-++--++++++++-+-+------+--+-+-+---+++-+--+---+

# todo - print a tree