-
Notifications
You must be signed in to change notification settings - Fork 0
/
7-2.py
35 lines (29 loc) · 973 Bytes
/
7-2.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from anytree.importer import DictImporter
from anytree.exporter import DotExporter
from anytree import RenderTree
from itertools import chain
a = open('7.txt').read().split('\n')[:-1]
a = [x.split(' -> ') for x in a]
a = [x[0].split() + [x[1].split(', ') if len(x)>1 else []] for x in a]
a = {x[0]: (int(x[1][1:-1]), x[2]) for x in a}
root = (set(x for x in a) - set(chain(*[a[x][1] for x in a]))).pop()
def get_total_weight(node):
children = a[node][1]
if len(children) == 0:
return a[node][0]
else:
return a[node][0] + sum([get_total_weight(x) for x in children])
def create_dict(node):
children = a[node][1]
if len(children) == 0:
return {
"a": a[node][0]
}
else:
return {
"a": (a[node][0], get_total_weight(node)),
"children": [create_dict(x) for x in children]
}
importer = DictImporter()
tree = importer.import_(create_dict(root))
print RenderTree(tree)