75 lines
1.4 KiB
Python
75 lines
1.4 KiB
Python
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])
|