day18_part1.py (8113B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 18 (part 1): Snailfish math 3 Snilfish math!""" 4 5 # Trying to see how simple I can keep this one: can a named tuple and global 6 # functions tackle this, or do I need a custom class with methods? 7 8 from typing import List, Tuple, Union, NamedTuple, cast 9 from math import ceil 10 from string import digits 11 from functools import reduce 12 from utils import get_puzzle_input 13 14 EXAMPLE_INPUT = \ 15 """[[[0,[5,8]],[[1,7],[9,6]]],[[4,[1,2]],[[1,4],2]]] 16 [[[5,[2,8]],4],[5,[[9,9],0]]] 17 [6,[[[6,2],[5,6]],[[7,6],[4,7]]]] 18 [[[6,[0,7]],[0,9]],[4,[9,[9,0]]]] 19 [[[7,[6,4]],[3,[1,3]]],[[[5,5],1],9]] 20 [[6,[[7,3],[3,2]]],[[[3,8],[5,7]],4]] 21 [[[[5,4],[7,7]],8],[[8,3],8]] 22 [[9,3],[[9,9],[6,[4,9]]]] 23 [[2,[[7,7],7]],[[5,8],[[9,3],[0,2]]]] 24 [[[[5,2],5],[8,[3,7]]],[[5,[7,5]],[4,4]]] 25 """ 26 27 # mypy can't handle recursive types yet (it's been an open issue since 2015 28 # (see issue #731), but that doesn't mean that we can't or shouldn't use them! 29 # Using TypedDict over NamedTuple because I want to use item assignment. 30 class MathNode(NamedTuple): 31 """Simple typed container for our snail math constructs""" 32 left: Union['MathNode', int] 33 right: Union['MathNode', int] 34 35 def parse_node(input_string: str) -> Union[MathNode, int, None]: 36 """Recursive parsing function""" 37 # Hey it's cool how we already did a problem where we kept track of 38 # brackets. Anyway, given a string that starts and ends with a bracket, use 39 # a counter to find the comma that splits it, and parse each side 40 # separately. If we're given an integer string, return that 41 if input_string in digits: 42 return int(input_string) 43 bracket_counter = 0 44 for idx, character in enumerate(input_string): 45 if character == "[": 46 bracket_counter += 1 47 elif character == "]": 48 bracket_counter -= 1 49 elif character == "," and bracket_counter == 1: 50 return MathNode(parse_node(input_string[1:idx]), 51 parse_node(input_string[(idx+1):-1])) 52 return None 53 54 def parse_input(input_string: str) -> List[MathNode]: 55 """Parses a list of snail math expressions, line-by-line""" 56 return [cast(MathNode, parse_node(line)) for line in \ 57 input_string.rstrip('\n').split('\n')] 58 59 def apply_to_leftmost(input_node: Union[MathNode, int], 60 value: Union[int, None]) -> Union[MathNode, int]: 61 """Apply the value to the leftmost value in the tree""" 62 if value is None: 63 return input_node 64 if isinstance(input_node, int): 65 return input_node + value 66 return MathNode(left = apply_to_leftmost(input_node.left, value), 67 right = input_node.right) 68 69 def apply_to_rightmost(input_node: Union[MathNode, int], 70 value: Union[int, None]) -> Union[MathNode, int]: 71 """Apply the value to the rightmost value in the tree""" 72 if value is None: 73 return input_node 74 if isinstance(input_node, int): 75 return input_node + value 76 return MathNode(left = input_node.left, 77 right = apply_to_rightmost(input_node.right, value)) 78 79 def explode_node(branch_node: Union[MathNode, int], 80 recursion_level: int = 0) -> \ 81 Tuple[Union[MathNode, int], bool, Union[int, None], Union[int, None]]: 82 """Recursively performs node "explosion". Returns a new (potentially) 83 exploded tree, a flag indicating whether explosions occurred, and values to 84 be applied to the left and right of the exploded node (if no values need to 85 be applied, these are None)""" 86 if isinstance(branch_node, int): 87 return (branch_node, False, None, None) 88 89 # Here is where explosion actually occurrs 90 if recursion_level == 4: 91 return (0, True, branch_node.left, branch_node.right) 92 93 # Handle the left child 94 left_child, did_explode, left_value, right_value = explode_node( 95 branch_node.left, recursion_level + 1 96 ) 97 if did_explode: 98 return (MathNode(left = left_child, 99 right = apply_to_leftmost(branch_node.right, 100 right_value)), 101 did_explode, left_value, None) 102 103 # Handle the right child 104 right_child, did_explode, left_value, right_value = explode_node( 105 branch_node.right, recursion_level + 1 106 ) 107 if did_explode: 108 return (MathNode(left = apply_to_rightmost(branch_node.left, 109 left_value), 110 right = right_child), 111 did_explode, None, right_value) 112 113 # If we made it here, this node and none of its children exploded, so it 114 # can return safely 115 return (branch_node, False, None, None) 116 117 def split_node(input_node: Union[MathNode, int]) -> \ 118 Tuple[Union[MathNode, int], bool]: 119 """Searching from left to right recursively, finds the first node to split (if any) and 120 splits it. Returns a tuple containing a (potentially) new tree structure 121 and a boolean indicating whether splitting occurred""" 122 if isinstance(input_node, int): 123 if input_node >= 10: 124 return (MathNode(left = input_node // 2, 125 right = ceil(input_node / 2)), 126 True) 127 # If it's an integer below 10 return it safely 128 return (input_node, False) 129 130 # Descend the tree, left-branch first 131 left_child, did_split = split_node(input_node.left) 132 if did_split: 133 return (MathNode(left = left_child, 134 right = input_node.right), 135 did_split) 136 137 right_child, did_split = split_node(input_node.right) 138 if did_split: 139 return (MathNode(left = input_node.left, 140 right = right_child), 141 did_split) 142 143 return (input_node, False) 144 145 def reduce_nodes(input_node: MathNode) -> MathNode: 146 """Launch the reduction steps""" 147 # From the puzzle text: "To reduce a snailfish number, you must repeatedly 148 # do the first action [between explode and split] that applies to the 149 # snailfish number". The text is a little ambiguous, but I'm currently 150 # understanding it to mean that no splits occur until all explosions are 151 # carried out, but that if a split creates an explosion condition, then all 152 # explosions must be carried out before we go back to splitting. This lends 153 # itself to a nested structure, with splitting occurring in the outer loop, 154 # and exploding happening in the inner loop. 155 output_node = input_node 156 reducing_done = False 157 while not reducing_done: 158 output_node, did_explode, _, _ = explode_node(output_node) 159 if not did_explode: 160 # If exploding didn't happen on this iteration of the loop, check 161 # to see if splitting is necessary 162 output_node, did_split = split_node(output_node) 163 if not did_split: 164 # If neither exploding NOR splitting happened, then we're done 165 # reducing. Otherwise, splitting occurred, and we need to check 166 # for exploding or splitting again 167 reducing_done = True 168 return output_node 169 170 def add_nodes(left_operand: MathNode, right_operand: MathNode) -> MathNode: 171 """Performs addition _and_ reduction on the given pair of MathNodes""" 172 return reduce_nodes(MathNode(left = left_operand, right = right_operand)) 173 174 def find_magnitude(input_node: Union[MathNode, int]) -> int: 175 """Find the magnitude of the final sum""" 176 if isinstance(input_node, int): 177 return input_node 178 return 3 * find_magnitude(input_node.left) + \ 179 2 * find_magnitude(input_node.right) 180 181 def solve_puzzle(input_string: str) -> int: 182 """Return the numeric solution to the puzzle""" 183 return find_magnitude( 184 reduce(add_nodes, 185 parse_input(input_string)) 186 ) 187 188 def main() -> None: 189 """Run when the file is called as a script""" 190 assert solve_puzzle(EXAMPLE_INPUT) == 4140 191 print("Magnitude of the final sum:", 192 solve_puzzle(get_puzzle_input(18))) 193 194 if __name__ == "__main__": 195 main()