aoc22/day25.py

converter = {
    '2' : 2,
    '1' : 1,
    '0' : 0,
    '-' : -1,
    '=' : -2
}

def maximum(n):
    maxi = 0
    for x in range(n + 1):
        maxi += 2*5**x
    return maxi

def decimalToSnafu(num):
    n = 0
    while True:
        if n == 0:
            if num <= 2:
                break
            n += 1
            continue
        if num <= maximum(n):
            break
        n += 1

    rest = num
    subtract = False
    if n > 0:
        Snafu = ['0' for x in range(n+1)]
        for digit in range(n, -1, -1):
            if rest > 0:
                count, rest = divmod(rest, 5**digit)
                if count != 2 and rest > maximum(digit - 1):
                    count += 1
                Snafu[-(digit+ 1)] = str(count)
            elif rest < 0:
                count, rest = divmod(abs(rest), 5**digit)
                if count != 2 and rest > maximum(digit - 1):
                    count += 1
                if count == 2:
                    Snafu[-(digit+ 1)] = '='
                elif count == 1:
                    Snafu[-(digit+ 1)] = '-'
                else:
                    Snafu[-(digit+ 1)] = '0'
            rest = num - SnafuToDecimal(''.join(Snafu)) 
        return(Snafu)
    else:
        return str(num)

def SnafuToDecimal(number):
    number = list(reversed(number))
    decimal = 0
    for x in range(len(number)):
        decimal += converter[number[x]] * 5 ** x
    return decimal


total = 0

with open('input25.txt', 'r') as f:
    numbers = f.read().splitlines(keepends=False)

for number in numbers:
    total += SnafuToDecimal(number)

print(total)
print(''.join(decimalToSnafu(total)))