Laskettelukeskus

#from collections import Counter
n = input()
lista = []
lista1 = []
lista2 = []
d = {}
try:
    for item in range(int(n)-1):
        a = input().split(" ")
        lista.append(a)
        lista1.append(a[0])
        lista2.append(a[1])
    lenlist = input().split(" ")
    if int(n) == 1:
        print(lenlist[0])
        exit()
    for l, item in enumerate(lenlist):
        lenlist[l] = int(item)
    leaves = set()
    leaves.update(*lista2)
    leaves -= set(lista1)
except Exception:
    print(":(")
def dfs(data, path, paths):
    datum = path[-1]              
    if datum in data:
        for val in data[datum]:
            new_path = path + [val]
            paths = dfs(data, new_path, paths)
    else:
        paths += [path]
    return paths

def enumerate_paths(graph):
    nodes = list(graph.keys())
    all_paths = []
    for node in nodes:
        node_paths = dfs(graph, [node], [])
        all_paths += node_paths
    return all_paths
   
for item in lista:
    if int(item[0]) in d:
        d[int(item[0])].append(int(item[1]))
    else:
        d[int(item[0])] = [int(item[1])]
paths = enumerate_paths(d)[:len(leaves)]
ma = max(len(x) for x in paths)
placeholder = 0
for i in range(ma):
    if max(lenlist) < 1:
        print(placeholder)
        exit()
    else: 
        for j, item in enumerate(paths):
            g = lenlist[item[-1]-1]
            placeholder += g
            for one in item:
                lenlist[one-1] -= g
                #print(lenlist)
                #print(item)
            #print("--------------------------")
            if lenlist[item[-1]-1] < 1:
                paths[j].pop(-1)
   

print(placeholder)

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 56, in <module>
    g = le...

Test 3

Group: 1, 2

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 56, in <module>
    g = le...

Test 4

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 5

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 6

Group: 1, 2

Verdict: ACCEPTED

Test 7

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 8

Group: 1, 2

Verdict: ACCEPTED

Test 9

Group: 2

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 47, in <module>
    paths...

Test 10

Group: 1, 2

Verdict: WRONG ANSWER

Test 11

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 12

Group: 1, 2

Verdict: WRONG ANSWER

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED