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):
    shortest_paths = {satama: (None, 0)}
    current_node = satama
    visited = set()
 
    while current_node != pelto:
        visited.add(current_node)
        destinations = []
        for i in graph[current_node]:
            destinations.append(i[0])
 
        weight_to_current_node = shortest_paths[current_node][1]
        for next_node in destinations:
            for i in graph[current_node]:
                if i[0] == next_node:
                    weight = i[1] + weight_to_current_node
            if next_node not in shortest_paths:
                shortest_paths[next_node] = (current_node, weight)
            else:
                current_shortest_weight = shortest_paths[next_node][1]
                if current_shortest_weight > weight:
                    shortest_paths[next_node] = (current_node, weight)
 
    
        next_destinations = {node: shortest_paths[node] for node in shortest_paths if node not in visited}
        current_node = min(next_destinations, key=lambda k: next_destinations[k][1])
    
    path = []
    path2 = 0
    while current_node is not None:
        path.append(current_node)
        next_node = shortest_paths[current_node][0]
        for i in graph[current_node]:
            if i[0] == next_node:
                path2 += i[1]
        current_node = next_node
    return path2    
 
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: ACCEPTED

Test 12

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 13

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 14

Group: 2

Verdict: ACCEPTED

Test 15

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 16

Group: 1, 2

Verdict: ACCEPTED

Test 17

Group: 1, 2

Verdict: ACCEPTED

Test 18

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

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: ACCEPTED

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