from functools import cmp_to_key import os input = open(os.path.dirname(__file__) + "/input.txt", "r").read().split('\n\n') def compare(a, b): typea = type(a) typeb = type(b) if typea == int and typeb == int: if a < b: return -1 elif a == b: return 0 else: return 1 elif typea == list and typeb == list: i = 0 while i < len(a) and i < len(b): res = compare(a[i], b[i]) if res == 1: return 1 if res == -1: return -1 i += 1 if len(a) == i and i < len(b): return -1 elif len(b) == i and i < len(a): return 1 else: return 0 elif (typea == int and typeb == list): return compare([a], b) elif (typeb == int and typea == list): return compare(a, [b]) packets = [] p1 = 0 for idx, i in enumerate(input): a, b = i.split('\n') a = eval(a) b = eval(b) packets.append(a) packets.append(b) if compare(a, b) == -1: p1 += idx+1 print("p1:", p1) packets.append([[2]]) packets.append([[6]]) packets = sorted(packets, key=cmp_to_key(lambda a, b: compare(a, b))) p2 = 1 for idx, p in enumerate(packets): # print(p) if p == [[2]] or p == [[6]]: p2 *= idx+1 print("p2:", p2)