from collections import defaultdict with open('data07.txt') as f: lines = f.readlines() class Tree: def __init__(self): self.items = {} def __getitem__(self, item): if item in self.items: return self.items[item] else: return Tree() def __setitem__(self, key, item): self.items[key] = item programs = defaultdict(dict) for line in lines: pieces = line.split() name = pieces[0] weight = int(pieces[1][1:-1]) prog = programs[name] prog.update({ 'name': name, 'weight': weight }) if len(pieces) > 3: # has children children = [c.strip(',') for c in pieces[3:]] prog['children'] = children for child in children: programs[child]['parent'] = name # part 1 ''' p = lines[-1].split()[0] while 'parent' in programs[p]: p = programs[p]['parent'] print(p) ''' # part 2 def weigh(p): if 'children' not in p: return p['weight'] else: total = p['weight'] for c in p['children']: child = programs[c] total += weigh(child) return total for prog in programs.values(): prog['total_weight'] = weigh(prog) current = 'qibuqqg' while 'children' in programs[current]: weights = defaultdict(list) for c in programs[current]['children']: child = programs[c] weight = child['total_weight'] weights[weight].append(c) if len(weights) == 1: break else: children = sorted(weights.values(), key=lambda x: len(x)) current = children[0][0] print(programs[current])