74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
import itertools
|
|
|
|
|
|
def parse_data() -> list[str]:
|
|
lines = []
|
|
with open("input.txt") as file:
|
|
for line in file:
|
|
lines.append(line.rstrip())
|
|
return lines
|
|
|
|
|
|
def solve_p1(program: list[str]) -> int:
|
|
bit_mask = ""
|
|
mem = {}
|
|
|
|
def apply_mask(_number: int, _mask: str) -> int:
|
|
res = _number
|
|
and_mask = int("".join(["1" if x == 'X' else "0" for x in _mask]), 2)
|
|
res = res & and_mask
|
|
or_mask = int("".join([x if x in "01" else "0" for x in _mask]), 2)
|
|
res = res | or_mask
|
|
return res
|
|
|
|
for line in program:
|
|
lvalue, rvalue = line.split(" = ")
|
|
if lvalue == "mask":
|
|
bit_mask = rvalue
|
|
elif lvalue[:3] == "mem":
|
|
address = lvalue[4:-1]
|
|
result = apply_mask(int(rvalue), bit_mask)
|
|
mem[address] = result
|
|
|
|
return sum(mem.values())
|
|
|
|
|
|
def solve_p2(program: list[str]) -> int:
|
|
bit_mask = ""
|
|
mem = {}
|
|
|
|
def apply_mask(_address: str, _mask: str) -> list[str]:
|
|
res = ['0' if x not in "01" else x for x in format(int(_address), "36b")]
|
|
for i in range(len(_mask)):
|
|
if _mask[i] == "1":
|
|
res[i] = "1"
|
|
elif _mask[i] == "X":
|
|
res[i] = "X"
|
|
return res
|
|
|
|
def resolve_floating(_address: list[str]) -> list[int]:
|
|
combinations = list(itertools.product(["0", "1"], repeat=_address.count("X")))
|
|
res = []
|
|
for combo in combinations:
|
|
res.append("".join(_address))
|
|
for c in combo:
|
|
res[-1] = res[-1].replace('X', c, 1)
|
|
return res
|
|
|
|
for line in program:
|
|
lvalue, rvalue = line.split(" = ")
|
|
if lvalue == "mask":
|
|
bit_mask = rvalue
|
|
elif lvalue[:3] == "mem":
|
|
address = lvalue[4:-1]
|
|
result = apply_mask(address, bit_mask)
|
|
floating_addresses = resolve_floating(result)
|
|
for floating_address in floating_addresses:
|
|
mem[floating_address] = int(rvalue)
|
|
|
|
return sum(mem.values())
|
|
|
|
|
|
print(solve_p1(parse_data()))
|
|
print(solve_p2(parse_data()))
|