83 lines
1.5 KiB
Python
83 lines
1.5 KiB
Python
from lib import *
|
|
|
|
EXAMPLE = """1=-0-2
|
|
12111
|
|
2=0=
|
|
21
|
|
2=01
|
|
111
|
|
20012
|
|
112
|
|
1=-1=
|
|
1-12
|
|
12
|
|
1=
|
|
122
|
|
"""
|
|
|
|
def snafu_to_dec(n: str) -> int:
|
|
r = 0
|
|
for i, c in enumerate(reversed(n)):
|
|
c = "=-012".index(c) - 2
|
|
r += c * 5**i
|
|
return r
|
|
|
|
def dec_to_snafu(n: int) -> str:
|
|
r = ""
|
|
while n != 0:
|
|
rem = n % 5
|
|
n //= 5
|
|
if rem <= 2:
|
|
r = str(rem) + r
|
|
else:
|
|
r = " =-"[rem] + r
|
|
n += 1
|
|
return r
|
|
|
|
def dec_to_snafu_cs(n: int) -> str:
|
|
import cpmpy as cp
|
|
import numpy as np
|
|
size = 20
|
|
xs = cp.intvar(-2, 2, shape=size, name="xs")
|
|
ns = np.array([5**i for i in range(size)])
|
|
m = cp.Model(cp.sum(ns * xs) == n)
|
|
m.solve()
|
|
|
|
value = xs.value()
|
|
assert value is not None
|
|
# print(value)
|
|
|
|
r = ""
|
|
for v in value:
|
|
if v >= 0:
|
|
r += str(v)
|
|
elif v == -1:
|
|
r += "-"
|
|
elif v == -2:
|
|
r += "="
|
|
else:
|
|
assert False
|
|
r = "".join(reversed(r))
|
|
r = r.lstrip("0")
|
|
return r
|
|
|
|
def solve(input: Input):
|
|
res = 0
|
|
for s in input.lines():
|
|
d = snafu_to_dec(s)
|
|
res += d
|
|
try:
|
|
assert dec_to_snafu_cs(res) == dec_to_snafu(res)
|
|
except ModuleNotFoundError:
|
|
print("Install 'cpmpy' to solve with constraint solver.")
|
|
return dec_to_snafu(res)
|
|
|
|
def main():
|
|
DAY_INPUT = "d25.txt"
|
|
print("Example 1:", solve(Input(EXAMPLE)))
|
|
print("Solution 1:", solve(Input(DAY_INPUT)))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|