Sadonkorjuu

import time
import sys
 
n = int(input())
k = input().split(" ")
k = list(map(int, k))
edges = []
for i in range(n - 1):
    a = input().split(" ")
    a = list(map(int, a))
    edges.append([a[0], a[1], a[2]])
 
start = time.time()
 
graph = {}
for edge in edges:
    a, b, c = edge[0], edge[1], edge[2]
 
    if a not in graph.keys():
        graph[a] = []
    if b not in graph.keys():
        graph[b] = []
 
    graph[a].append([b, c])
    graph[b].append([a, c])
 
def findSatamas(pelto):
    satamas = []
    visited = [pelto]
    queue = [pelto]
 
    while queue:
        s = queue.pop(0)
        for neighbour in graph[s]:
            if neighbour[0] not in visited:
                if k[neighbour[0] - 1] == 0:
                    satamas.append(neighbour[0])
                else:
                    queue.append(neighbour[0])
                visited.append(neighbour[0])
    return satamas
 
def dist(satama, pelto):
    unvisited = list(graph.keys())
 
    shortest = {}
    previous = {}
 
    for node in unvisited:
        shortest[node] = sys.maxsize
 
    shortest[pelto] = 0
    shortestSatama = sys.maxsize
 
    while unvisited:
        min = None
        for node in unvisited:
            if min == None:
                min = node
            elif shortest[node] < shortest[min]:
                min = node
 
        neighbors = []
        values = []
        for i in graph[min]:
            neighbors.append(i[0])
            values.append(i[1])
        for neighbor in neighbors:
            value = shortest[min] + values[neighbors.index(neighbor)]
            if value > shortestSatama:
                continue
            if value < shortest[neighbor]:
                shortest[neighbor] = value
                if neighbor == satama:
                    shortestSatama = value
                previous[neighbor] = min
 
        unvisited.remove(min)
    
    return shortestSatama
 
ans = 0
while 1 in k:
    pelto = k.index(1)
    satamas = findSatamas(pelto + 1)
    distances = []
    for satama in satamas:
        distances.append(dist(satama, pelto + 1))
    ans += min(distances)
    k[pelto] = 3
 
print(ans)

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

Test 7

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 8

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

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

Test 16

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 17

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 18

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 19

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 20

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

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

Test 25

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 26

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 27

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 28

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 29

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 30

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 31

Group: 2

Verdict: TIME LIMIT EXCEEDED