aoc/2022/25/solution.py

#!/usr/bin/python3
import sys
import math

snafu_to_base10 = { '=': -2, '-': -1, '0': 0,'1': 1,'2': 2, }
base5_to_snafu = { -2: '=', -1: '-', 0: '0', 1: '1', 2: '2', }

def construct_number(target, digits, num, snafu):
    if num < 0:
        return snafu, num
    if digits == 0:
        return snafu, num
    if target - num > 5 * 5**digits:
        print(target - num, 2 * 5**digits)
        print(target, digits, num, snafu)
        return snafu, num
    for i in [ '1', '2', '0', '-', '=' ]:
        diff = 5**digits * snafu_to_base10[i]
        if digits % 9 == 0: print(target, digits, num, snafu)
        curnum = num + diff
        new_snafu = snafu + i

        new_snafu, result = construct_number(target, digits - 1, curnum, new_snafu)
        if result == target:
            return new_snafu, result
    return snafu, -1

def to_base_b(num, b):
    num_n = 0
    i = 0
    while num > 0:
        num, rem = divmod(num, b)
        num_n += 10**i * rem
        i += 1

    return num_n

def solution():
    numsum = 0
    for line in sys.stdin:
        num = sum([snafu_to_base10[digit] * 5**i for i, digit in enumerate(line.strip()[::-1])])
        numsum += num
        print(num)
    print(numsum)

    #snafu, _ = construct_number(numsum, len(str(numsum)) + 5, 0, '')
    #snafu = snafu.lstrip('0')
    #print(f'the sum {numsum} has snafu number {snafu}')

    result = ''
    carry = 0
    for digit in str(to_base_b(numsum, 5))[::-1]:
        digint = int(digit) + carry
        if digint > 2:
            result += base5_to_snafu[digint - 5]
            carry = 1
        else:
            result += base5_to_snafu[digint]
            carry = 0
    print(result[::-1])

if sys.argv[1] in '1':
    solution()
else:
    solution()