Sadonkorjuu

#import math
#import numpy
import math
CityCount = int(input())
CityStates = input().split()
RoadList = []
 
#Build a nested array of roads
for n in range(0,CityCount-1):
    temp = input().split()
    RoadList.append(temp)
 
#Convert to numpy array
#RoadMatrix = numpy.array(RoadList)
 
RoadDict = {}
 
#Cycle through roads to find ones that connect to a given node, then build a dictionary which indexes each connected road
for a in range(0,CityCount):
    tempList = []
    #Roads
    for b in range(0, CityCount-1):
        if(int(RoadList[b][0]) == a+1):
           tempList.append([int(RoadList[b][1])-1,int(RoadList[b][2])])
        elif(int(RoadList[b][1]) == a+1):
           tempList.append([int(RoadList[b][0])-1,int(RoadList[b][2])])
    
    #Cycle them through again, rearraging them in ascending weight
    for b in range(0, len(tempList)):
        tempList.sort(key= lambda x: x[1])
    RoadDict[a] = tempList
 
#Actually find something
currentNode = 0
nodeDistances = [0]+[math.inf]*(CityCount-1)
totalDistances = 0
 
unvisited_nodes = [x for x in range(CityCount)]
shortest_path = {}
previous_nodes = {}
start_node = 0
 
def connection(node, target):
    node = node
    target = target
    for con in RoadDict[node]:
        if(int(con[0]) == target):
            return int(con[1])
    return 1
 
            
 
def Distance(nodeList):
    totalDistance = 0
    for nodeIndex in range(0,len(nodeList)-1):
        totalDistance += connection(nodeList[nodeIndex], nodeList[nodeIndex+1])
    return totalDistance
 
for i in range(0, CityCount):
    #Checking if is a producer node, not a port node
    totalDistance = 0
    subPath = 0
    pathDistances = {}
    if(CityStates[i] == "1"):
        visitingNode = i
        nodeCache = i
        #Loop stays true until port is found
        portfound = False
        visitedNodes = [False]*(CityCount)
        nodeHistory = [i]
        smallestWeight = math.inf
        smallestWeightNode = 0
        while(not portfound):
            nothing = True 
            #Cycle through neighbors
            for neighbor in RoadDict[visitingNode]:
                unvisitedNodes = [x[0] for x in RoadDict[visitingNode]]
                #If port is found, end the search
                if(CityStates[int(neighbor[0])] =="0"):
                    #Port found!
                    nodeHistory.append(int(neighbor[0]))
                    portfound, nothing = True, False
 
                elif(int(neighbor[1]) < smallestWeight and not visitedNodes[int(neighbor[0])]):
                    smallestWeight = int(neighbor[1])
                    smallestWeightNode = int(neighbor[0])
                    nothing = False
            if(not portfound and not nothing):
                smallestWeight = math.inf
                visitedNodes[visitingNode] = True
                nodeHistory.append(int(smallestWeightNode)) 
                visitingNode = smallestWeightNode
            
            elif(nothing):
                visitedNodes[nodeHistory[-2]] = False
                visitedNodes[nodeHistory[-1]] = True
                smallestWeight = math.inf
                visitingNode = nodeHistory[-1]
                nodeHistory = nodeHistory[:-2]
        totalDistances += Distance(nodeHistory)            
                
 
print(totalDistances)

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: WRONG ANSWER

Test 6

Group: 1, 2

Verdict: TIME LIMIT EXCEEDED

Test 7

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 8

Group: 1, 2

Verdict: WRONG ANSWER

Test 9

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 10

Group: 1, 2

Verdict: WRONG ANSWER

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

Test 16

Group: 1, 2

Verdict: WRONG ANSWER

Test 17

Group: 1, 2

Verdict: WRONG ANSWER

Test 18

Group: 1, 2

Verdict: WRONG ANSWER

Test 19

Group: 1, 2

Verdict: WRONG ANSWER

Test 20

Group: 1, 2

Verdict: WRONG ANSWER

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: WRONG ANSWER

Test 25

Group: 2

Verdict: TIME LIMIT EXCEEDED

Test 26

Group: 1, 2

Verdict: WRONG ANSWER

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