Sadonkorjuu

import sys
 
n = int(input())
k = list(map(int, input().split(" ")))
edges = {}
for i in range(n - 1):
    a = list(map(int, input().split(" ")))
    if a[0] not in edges.keys():
        edges[a[0]] = []
    if a[1] not in edges.keys():
        edges[a[1]] = []
    edges[a[0]].append([a[1], a[2]])
    edges[a[1]].append([a[0], a[2]])
if not edges: edges[1] = [[1, 0]]
 
shortest = {}
satamas = {}
ends = []
for i in range(len(k)):
    if len(edges[i + 1]) == 1: ends.append(i + 1)
    if k[i] == 0:
        satamas[i + 1] = True
        shortest[i + 1] = 0
    else:
        satamas[i + 1] = False
        shortest[i + 1] = sys.maxsize
 
def syvyyshaku(point):
    if used[point - 1]: return
    used[point - 1] = True
    for i in edges[point]:
        if used[i[0] - 1]: continue
        if satamas[i[0]]: w[i[0]] = 0
        else: w[i[0]] = w[point] + i[1]
        if w[i[0]] < shortest[i[0]]: shortest[i[0]] = w[i[0]]
        syvyyshaku(i[0])

for i in ends:
    used = [False] * n
    w = {}
    w[i] = shortest[i]
    syvyyshaku(i)

print(sum(shortest.values()))

Test 1

Group: 1, 2

Verdict: ACCEPTED

Test 2

Group: 1, 2

Verdict: ACCEPTED

Test 3

Group: 1, 2

Verdict: ACCEPTED

Test 4

Group: 1, 2

Verdict: ACCEPTED

Test 5

Group: 1, 2

Verdict: ACCEPTED

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: TIME LIMIT EXCEEDED

Test 10

Group: 1, 2

Verdict: ACCEPTED

Test 11

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 12

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 14

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 15

Group: 1, 2

Verdict: RUNTIME ERROR

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

Test 16

Group: 1, 2

Verdict: RUNTIME ERROR

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

Test 17

Group: 1, 2

Verdict: RUNTIME ERROR

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

Test 18

Group: 1, 2

Verdict: ACCEPTED

Test 19

Group: 1, 2

Verdict: ACCEPTED

Test 20

Group: 1, 2

Verdict: ACCEPTED

Test 21

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 22

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 23

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 24

Group: 1, 2

Verdict: ACCEPTED

Test 25

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 26

Group: 1, 2

Verdict: ACCEPTED

Test 27

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 28

Group: 1, 2

Verdict: ACCEPTED

Test 29

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 30

Group: 1, 2

Verdict: ACCEPTED

Test 31

Group: 2

Verdict: TIME LIMIT EXCEEDED