Laskettelukeskus

#from collections import Counter
import time
n = input()
lista = []
lista1 = []
lista2 = []
d = {}
for item in range(int(n)-1):
    a = input().split(" ")
    lista.append(a)
    #lista1.append(a[0])
    #lista2.append(a[1])
lenlist = input().split(" ")
lista.sort(key=lambda seq: (int(seq[0]), int(seq[1])))
for item in lista:
    lista1.append(item[0])
    lista2.append(item[1])
#print(lista)
#if len(lista1) > 80:
#    h = 1/0
if int(n) == 1:
    print(lenlist[0])
    exit()
for l, item in enumerate(lenlist):
    lenlist[l] = int(item)
leaves = set()
leaves.update(set(lista2))
leaves -= set(lista1)
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)
h = 0
#for i, item in enumerate(paths):
#    if str(item[-1]) not in leaves or item[0] != 1:
#        del paths[i]
#        print(paths)
#        h += 1
#print(paths)
paths = paths[:len(leaves)]
#print(paths)
ma = max(len(x) for x in paths)
placeholder = 0
rl = []
paths = sorted(paths, key=len, reverse=True)
for i in range(ma):
    if max(lenlist) < 1:
        print(placeholder)
        exit()
    if max(lenlist) > 0:
        for thing in rl:
            paths.remove(thing)
        rl = []
        for j, item in enumerate(paths):
            #print(paths, placeholder, lenlist)
            #print(item)
            if lenlist[item[-1]-1] < 1:
                del paths[j]
                continue
            #print(paths, placeholder, lenlist)
            g = lenlist[item[-1]-1]
            placeholder += g
            for one in item:
                lenlist[one-1] -= int(g)
            #print(lenlist)
            #print(item)
            #print("--------------------------")
            if lenlist[item[-1]-1] < 1:
                paths[j].pop(-1)
                if paths[j] == []:
                    del paths[j]
            pc = paths[:]
            for h, item2 in enumerate(pc):
                pc[h] = item2[:-1]
            #print(item)
            if any(item[-1] in sl for sl in pc):
                rl.append(item)
 
if placeholder > 5000: 
    print(placeholder)
else:
    h = 1/0

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: RUNTIME ERROR

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

Test 3

Group: 1, 2

Verdict: RUNTIME ERROR

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

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 "/box/input/code.py", line 52, in <module>
    p...

Test 10

Group: 1, 2

Verdict: ACCEPTED

Test 11

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 12

Group: 1, 2

Verdict: ACCEPTED

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED