Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from utils import *
- class Node:
- def __init__(self, parent, data, left=None, right=None):
- self.parent = parent
- self.data = data
- self.left = left
- self.right = right
- def __str__(self):
- if self.data is not None:
- return str(self.data)
- else:
- return f"[{self.left}, {self.right}]"
- def parse(raw):
- # convert nested array into binary tree
- def convert(data, parent=None):
- if type(data) == int:
- return Node(parent, data)
- else:
- node = Node(parent, None)
- left = convert(data[0], node)
- right = convert(data[1], node)
- node.left = left
- node.right = right
- return node
- return convert(raw)
- def runref(data):
- val = data
- while True:
- try:
- linear = linearlize(val)
- red1(linear, val, 0)
- red2(linear, val, 0)
- return val
- # Forgive me for I have sinned
- except Exception as e:
- # print(e.with_traceback())
- pass
- def linearlize(cur):
- if cur.data is None:
- return [cur, *linearlize(cur.left), *linearlize(cur.right)]
- else:
- return [cur]
- def red1(linear: list, node, depth):
- if node.data is None and depth > 3 and node.left.data is not None and node.right.data is not None:
- l = node.left.data
- r = node.right.data
- index = linear.index(node)
- for i in range(index, -1, -1):
- if linear[i].data is not None:
- linear[i].data += l
- break
- for i in range(index + 3, len(linear)):
- if linear[i].data is not None:
- linear[i].data += r
- break
- # Zero out this node
- node.data = 0
- node.left = None
- node.right = None
- raise Exception()
- elif node.data is None:
- red1(linear, node.left, depth + 1)
- red1(linear, node.right, depth + 1)
- def red2(linear: list, node, depth):
- if node.data is not None:
- # Check if split is needed
- if node.data > 9:
- node.left = Node(node, math.floor(node.data / 2))
- node.right = Node(node, math.ceil(node.data / 2))
- node.data = None
- raise Exception()
- elif node.data is None:
- red2(linear, node.left, depth + 1)
- red2(linear, node.right, depth + 1)
- def mag(node):
- if node.data is None:
- return 3*mag(node.left) + 2*mag(node.right)
- else:
- return node.data
- def run():
- file = open(sys.argv[1] if len(sys.argv) > 1 else 'input.txt', "r")
- data = file.read()
- lines = data.splitlines()
- m = 0
- for _l in lines:
- for _r in lines:
- if _l != _r:
- l = parse(eval(_l))
- r = parse(eval(_r))
- newv = Node(None, None)
- newv.left = l
- newv.right = r
- l.parent = newv
- r.parent = newv
- v = runref(newv)
- m = max(m, mag(v))
- print(m)
- if __name__ == "__main__":
- run()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement