adventofcode2022/day25/sol.py

def eval(num):
    digits = ["=-012".index(c)-2 for c in num]
    n = 0
    for x in digits:
        n = n*5 + x
    return n

def tosnafu(n):
    orig = n
    digits = []
    carry = 0
    while n != 0:
        n, d = divmod(n, 5)
        carry, d = divmod(d+carry, 5)
        if d > 2:
            carry += 1
            d -= 5
        digits.append(d)
    if carry:
        digits.append(carry)
    digits.reverse()

    snafu = "".join("=-012"[d+2] for d in digits)
    #print(n, digits, snafu)
    assert eval(snafu) == orig
    return snafu

t = 0
for line in open("input"):
    n = eval(line.strip())
    print(n)
    t += n

print(t)
print(tosnafu(t))