Advertisement
Guest User

Untitled

a guest
Dec 18th, 2021
374
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 5.26 KB | None | 0 0
  1. #!/usr/bin/python3
  2.  
  3. import ast
  4. import copy
  5. import math
  6. import sys
  7.  
  8. class SnailfishNode:
  9. def __init__(self, parent=None, left=None, right=None, value=None):
  10. self.value = value
  11. self.left = None
  12. if left is not None:
  13. self.left = left
  14. self.left.parent = self
  15. self.right = None
  16. if right is not None:
  17. self.right = right
  18. self.right.parent = self
  19. self.parent = parent
  20.  
  21. def __str__(self):
  22. if self.value is not None:
  23. return str(self.value)
  24. else:
  25. return '[' + str(self.left) + ',' + str(self.right) + ']'
  26.  
  27. def explode_rec(self, sfn, level):
  28. if sfn.value is None:
  29. if level < 4:
  30. (exploded, node) = self.explode_rec(sfn.left, level+1)
  31. if exploded:
  32. return (exploded, node)
  33. (exploded, node) = self.explode_rec(sfn.right, level+1)
  34. if exploded:
  35. return (exploded, node)
  36. return (False, None)
  37. else:
  38. return (True, sfn)
  39. else:
  40. return (False, None)
  41.  
  42. def find_explode_if_needed(self):
  43. return self.explode_rec(self, 0)
  44.  
  45. def add_to_left(self, node, number):
  46. while True:
  47. if node.parent is None:
  48. return
  49. if node == node.parent.right:
  50. self.add_rightmost(node.parent.left, number)
  51. return
  52. else:
  53. node = node.parent
  54.  
  55. def add_rightmost(self, node, number):
  56. while node.right:
  57. node = node.right
  58. node.value += number
  59.  
  60. def add_to_right(self, node, number):
  61. while True:
  62. if node.parent is None:
  63. return
  64. if node == node.parent.left:
  65. self.add_leftmost(node.parent.right, number)
  66. return
  67. else:
  68. node = node.parent
  69.  
  70. def add_leftmost(self, node, number):
  71. while node.left:
  72. node = node.left
  73. node.value += number
  74.  
  75. def explode(self, node):
  76. left = node.left.value
  77. right = node.right.value
  78. node.left = None
  79. node.right = None
  80. node.value = 0
  81. self.add_to_left(node, left)
  82. self.add_to_right(node, right)
  83.  
  84. def split_rec(self, node):
  85. if node.value is not None:
  86. if node.value >= 10:
  87. return (True, node)
  88. else:
  89. return (False, node)
  90. else:
  91. split, split_node = self.split_rec(node.left)
  92. if split:
  93. return (split, split_node)
  94. split, split_node = self.split_rec(node.right)
  95. if split:
  96. return (split, split_node)
  97. return (False, None)
  98.  
  99. def find_split_if_needed(self):
  100. return self.split_rec(self)
  101.  
  102. def split(self, node):
  103. node.left = SnailfishNode(value=int(math.floor(node.value / 2)), parent=node)
  104. node.right = SnailfishNode(value=int(math.ceil(node.value / 2)), parent=node)
  105. node.value = None
  106.  
  107. def reduce(self):
  108. while True:
  109. exploded, node = self.find_explode_if_needed()
  110. if exploded:
  111. self.explode(node)
  112. continue
  113. split, node = self.find_split_if_needed()
  114. if split:
  115. self.split(node)
  116. continue
  117. return
  118.  
  119. def magnitude(self):
  120. def magnitude_rec(node):
  121. if node.value is not None:
  122. return node.value
  123. else:
  124. return magnitude_rec(node.left) * 3 + magnitude_rec(node.right) * 2
  125.  
  126. return magnitude_rec(self)
  127.  
  128. def parse_rec(sfn, data):
  129. if isinstance(data, list):
  130. sfn.left = SnailfishNode()
  131. parse_rec(sfn.left, data[0])
  132. sfn.left.parent = sfn
  133. sfn.right = SnailfishNode()
  134. parse_rec(sfn.right, data[1])
  135. sfn.right.parent = sfn
  136. else:
  137. sfn.value = data
  138.  
  139. def parse_snailfish(data):
  140. sfn_data = ast.literal_eval(data)
  141. sfn = SnailfishNode()
  142. parse_rec(sfn, sfn_data)
  143. return sfn
  144.  
  145. def sum_all(lines):
  146. sfn_sum = None
  147. for line in lines:
  148. sfn = parse_snailfish(line)
  149. sfn.reduce()
  150. if sfn_sum is None:
  151. sfn_sum = sfn
  152. else:
  153. sfn_sum = SnailfishNode(left=sfn_sum, right=sfn)
  154. sfn_sum.reduce()
  155. return sfn_sum.magnitude()
  156.  
  157. def sum_two(lines):
  158. fish = []
  159. for line in lines:
  160. sfn = parse_snailfish(line)
  161. sfn.reduce()
  162. fish.append(sfn)
  163.  
  164. max_magnitude = 0
  165. for i in range(len(fish)):
  166. for j in range(len(fish)):
  167. if i != j:
  168. sfn_sum = SnailfishNode(left=copy.deepcopy(fish[i]), right=copy.deepcopy(fish[j]))
  169. sfn_sum.reduce()
  170. magnitude = sfn_sum.magnitude()
  171. if magnitude > max_magnitude:
  172. max_magnitude = magnitude
  173. return max_magnitude
  174.  
  175. def main():
  176. lines = [line.strip() for line in sys.stdin.readlines()]
  177. # Part1
  178. print(sum_all(lines))
  179. # Part2
  180. print(sum_two(lines))
  181.  
  182. if __name__ == "__main__":
  183. main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement