233 lines
7.6 KiB
Python
233 lines
7.6 KiB
Python
import json
|
|
import re
|
|
from itertools import combinations
|
|
from math import ceil
|
|
from tools.aoc import AOCDay
|
|
from typing import Any, List, Union
|
|
|
|
re_pair = re.compile(r"\[(\d+),(\d+)\]")
|
|
re_lnum = re.compile(r".*[^\d](\d+)")
|
|
re_rnum = re.compile(r"(\d+)")
|
|
re_bignum = re.compile(r"(\d\d+)")
|
|
|
|
|
|
class Snailfish:
|
|
def __init__(self, fish_str: str):
|
|
self.pairs = fish_str
|
|
while self.reduce():
|
|
pass
|
|
|
|
def __add__(self, other: 'Snailfish'):
|
|
return Snailfish("[%s,%s]" % (self.pairs, other.pairs))
|
|
|
|
def reduce(self):
|
|
open_count = 0
|
|
for i, c in enumerate(self.pairs):
|
|
if c == "]":
|
|
open_count -= 1
|
|
elif c == "[":
|
|
open_count += 1
|
|
if open_count == 5:
|
|
explode_pair_match = re_pair.search(self.pairs, i)
|
|
index_l, index_r = explode_pair_match.span()
|
|
explode_left, explode_right = map(int, explode_pair_match.groups())
|
|
lnum_old = rnum_old = lnum_new = rnum_new = ""
|
|
lnum_index = index_l
|
|
rnum_index = index_r
|
|
|
|
if left_num := re_lnum.search(self.pairs[:index_l]):
|
|
lnum_old = left_num.groups()[0]
|
|
lnum_new = str(int(lnum_old) + explode_left)
|
|
lnum_index = left_num.span()[1] - len(lnum_old)
|
|
|
|
if right_num := re_rnum.search(self.pairs, index_r):
|
|
rnum_old = right_num.groups()[0]
|
|
rnum_new = str(int(rnum_old) + explode_right)
|
|
rnum_index = right_num.span()[0]
|
|
|
|
self.pairs = "".join([
|
|
self.pairs[:lnum_index],
|
|
lnum_new,
|
|
self.pairs[lnum_index + len(lnum_old):index_l],
|
|
"0",
|
|
self.pairs[index_r:rnum_index],
|
|
rnum_new,
|
|
self.pairs[rnum_index + len(rnum_old):]
|
|
])
|
|
return True
|
|
|
|
if big_num_match := re_bignum.search(self.pairs):
|
|
num = int(big_num_match.group())
|
|
self.pairs = self.pairs.replace(big_num_match.group(), "[%d,%d]" % (num // 2, int(ceil(num / 2))), 1)
|
|
return True
|
|
|
|
return False
|
|
|
|
def getMagnitude(self) -> int:
|
|
pairs = self.pairs
|
|
while match := re_pair.findall(pairs):
|
|
for num1, num2 in match:
|
|
pairs = pairs.replace("[%s,%s]" % (num1, num2), str(int(num1) * 3 + int(num2) * 2))
|
|
|
|
return int(pairs)
|
|
|
|
|
|
class BinarySnailfish:
|
|
def __init__(self, value: Union[int, List] = None, depth: int = 0, parent: 'BinarySnailfish' = None):
|
|
self.fours = []
|
|
self.depth = depth
|
|
self.parent = parent
|
|
if not isinstance(value, list):
|
|
self.value = value
|
|
self.left = None
|
|
self.right = None
|
|
else:
|
|
self.value = None
|
|
if isinstance(value[0], BinarySnailfish):
|
|
self.left = value[0]
|
|
self.left.setDepth(self.depth + 1)
|
|
self.left.parent = self
|
|
else:
|
|
self.left = BinarySnailfish(value[0], parent=self, depth=self.depth + 1)
|
|
|
|
if isinstance(value[1], BinarySnailfish):
|
|
self.right = value[1]
|
|
self.right.setDepth(self.depth + 1)
|
|
self.right.parent = self
|
|
else:
|
|
self.right = BinarySnailfish(value[1], parent=self, depth=self.depth + 1)
|
|
|
|
def setDepth(self, depth: int):
|
|
self.depth = depth
|
|
if self.value is not None:
|
|
return
|
|
else:
|
|
if self.depth == 4:
|
|
self.getRoot().fours.append(self)
|
|
self.left.setDepth(depth + 1)
|
|
self.right.setDepth(depth + 1)
|
|
|
|
def getRoot(self) -> 'BinarySnailfish':
|
|
root = self
|
|
while root.parent:
|
|
root = root.parent
|
|
|
|
return root
|
|
|
|
def getLeftNumber(self) -> Union[None, 'BinarySnailfish']:
|
|
parent = self.parent
|
|
if self == self.parent.left:
|
|
while parent.parent and parent == parent.parent.left:
|
|
parent = parent.parent
|
|
parent = parent.left
|
|
else:
|
|
while parent.parent and parent == parent.parent.right:
|
|
parent = parent.parent
|
|
parent = parent.right
|
|
while parent.right.value is None:
|
|
parent = parent.right
|
|
return parent
|
|
|
|
leftnum = None
|
|
for current in self.getRoot().traverse_preorder():
|
|
if current == self:
|
|
return leftnum
|
|
if current.value is not None:
|
|
leftnum = current
|
|
|
|
def getRightNumber(self) -> Union[None, 'BinarySnailfish']:
|
|
returnNext = 0
|
|
for current in self.getRoot().traverse_preorder():
|
|
if returnNext >= 3 and current.value is not None:
|
|
return current
|
|
if returnNext > 0:
|
|
returnNext += 1
|
|
if current == self:
|
|
returnNext += 1
|
|
|
|
def split(self):
|
|
self.left = BinarySnailfish(self.value // 2, depth=self.depth + 1, parent=self)
|
|
self.right = BinarySnailfish(int(ceil(self.value / 2)), depth=self.depth + 1, parent=self)
|
|
self.value = None
|
|
|
|
def explode(self):
|
|
left_num = self.getLeftNumber()
|
|
if left_num is not None and left_num.value is not None:
|
|
left_num.value += self.left.value
|
|
|
|
right_num = self.getRightNumber()
|
|
if right_num is not None and right_num.value is not None:
|
|
right_num.value += self.right.value
|
|
|
|
self.left = self.right = None
|
|
self.value = 0
|
|
|
|
def reduce(self):
|
|
if self.value is not None:
|
|
return
|
|
found = True
|
|
while found:
|
|
found = False
|
|
for n in self.traverse_preorder():
|
|
if n.value is None and n.depth == 4:
|
|
n.explode()
|
|
found = True
|
|
break
|
|
|
|
if found:
|
|
continue
|
|
|
|
for n in self.traverse_preorder():
|
|
if n.value is not None and n.value > 9:
|
|
n.split()
|
|
found = True
|
|
break
|
|
|
|
def traverse_preorder(self):
|
|
yield self
|
|
|
|
if self.left:
|
|
for x in self.left.traverse_preorder():
|
|
yield x
|
|
|
|
if self.right:
|
|
for x in self.right.traverse_preorder():
|
|
yield x
|
|
|
|
def __add__(self, other):
|
|
sum_fish = BinarySnailfish([self, other])
|
|
sum_fish.reduce()
|
|
return sum_fish
|
|
|
|
def __str__(self):
|
|
if self.value is None:
|
|
return "[%s,%s]" % (str(self.left), str(self.right))
|
|
else:
|
|
return str(self.value)
|
|
|
|
def getMagnitude(self):
|
|
if self.value is not None:
|
|
return self.value
|
|
else:
|
|
return self.left.getMagnitude() * 3 + self.right.getMagnitude() * 2
|
|
|
|
|
|
class Day(AOCDay):
|
|
test_solutions_p1 = [4140, 4417]
|
|
test_solutions_p2 = [3993, 4796]
|
|
|
|
def part1(self) -> Any:
|
|
snailfishes = [Snailfish(x) for x in self.getInput()]
|
|
#snailfishes = [BinarySnailfish(json.loads(x)) for x in self.getInput()]
|
|
return sum(snailfishes[1:], start=snailfishes[0]).getMagnitude()
|
|
|
|
def part2(self) -> Any:
|
|
snailfishes = [Snailfish(x) for x in self.getInput()]
|
|
max_mag = 0
|
|
for a, b in combinations(snailfishes, 2):
|
|
sub_mag_a = (a + b).getMagnitude()
|
|
sub_mag_b = (b + a).getMagnitude()
|
|
max_mag = max(max_mag, sub_mag_a, sub_mag_b)
|
|
|
|
return max_mag
|