Sadonkorjuu

import math

class FibonacciHeap:

    # internal node class
    class Node:
        def __init__(self, key, value):
            self.key = key
            self.value = value
            self.parent = self.child = self.left = self.right = None
            self.degree = 0
            self.mark = False

    # function to iterate through a doubly linked list
    def iterate(self, head):
        node = stop = head
        flag = False
        while True:
            if node == stop and flag is True:
                break
            elif node == stop:
                flag = True
            yield node
            node = node.right

    # pointer to the head and minimum node in the root list
    root_list, min_node = None, None

    # maintain total node count in full fibonacci heap
    total_nodes = 0

    # return min node in O(1) time
    def find_min(self):
        return self.min_node

    # extract (delete) the min node from the heap in O(log n) time
    # amortized cost analysis can be found here (http://bit.ly/1ow1Clm)
    def extract_min(self):
        z = self.min_node
        if z is not None:
            if z.child is not None:
                # attach child nodes to root list
                children = [x for x in self.iterate(z.child)]
                for i in range(0, len(children)):
                    self.merge_with_root_list(children[i])
                    children[i].parent = None
            self.remove_from_root_list(z)
            # set new min node in heap
            if z == z.right:
                self.min_node = self.root_list = None
            else:
                self.min_node = z.right
                self.consolidate()
            self.total_nodes -= 1
        return z

    # insert new node into the unordered root list in O(1) time
    def insert(self, key, value=None):
        n = self.Node(key, value)
        n.left = n.right = n
        self.merge_with_root_list(n)
        if self.min_node is None or n.key < self.min_node.key:
            self.min_node = n
        self.total_nodes += 1
        return n

    # modify the key of some node in the heap in O(1) time
    def decrease_key(self, x, k):
        if k > x.key:
            return None
        x.key = k
        y = x.parent
        if y is not None and x.key < y.key:
            self.cut(x, y)
            self.cascading_cut(y)
        if x.key < self.min_node.key:
            self.min_node = x

    # merge two fibonacci heaps in O(1) time by concatenating the root lists
    # the root of the new root list becomes equal to the first list and the second
    # list is simply appended to the end (then the proper min node is determined)
    def merge(self, h2):
        H = FibonacciHeap()
        H.root_list, H.min_node = self.root_list, self.min_node
        # fix pointers when merging the two heaps
        last = h2.root_list.left
        h2.root_list.left = H.root_list.left
        H.root_list.left.right = h2.root_list
        H.root_list.left = last
        H.root_list.left.right = H.root_list
        # update min node if needed
        if h2.min_node.key < H.min_node.key:
            H.min_node = h2.min_node
        # update total nodes
        H.total_nodes = self.total_nodes + h2.total_nodes
        return H

    # if a child node becomes smaller than its parent node we
    # cut this child node off and bring it up to the root list
    def cut(self, x, y):
        self.remove_from_child_list(y, x)
        y.degree -= 1
        self.merge_with_root_list(x)
        x.parent = None
        x.mark = False

    # cascading cut of parent node to obtain good time bounds
    def cascading_cut(self, y):
        z = y.parent
        if z is not None:
            if y.mark is False:
                y.mark = True
            else:
                self.cut(y, z)
                self.cascading_cut(z)

    # combine root nodes of equal degree to consolidate the heap
    # by creating a list of unordered binomial trees
    def consolidate(self):
        A = [None] * self.total_nodes
        nodes = [w for w in self.iterate(self.root_list)]
        for w in range(0, len(nodes)):
            x = nodes[w]
            d = x.degree
            while A[d] != None:
                y = A[d]
                if x.key > y.key:
                    temp = x
                    x, y = y, temp
                self.heap_link(y, x)
                A[d] = None
                d += 1
            A[d] = x
        # find new min node - no need to reconstruct new root list below
        # because root list was iteratively changing as we were moving
        # nodes around in the above loop
        for i in range(0, len(A)):
            if A[i] is not None:
                if A[i].key < self.min_node.key:
                    self.min_node = A[i]

    # actual linking of one node to another in the root list
    # while also updating the child linked list
    def heap_link(self, y, x):
        self.remove_from_root_list(y)
        y.left = y.right = y
        self.merge_with_child_list(x, y)
        x.degree += 1
        y.parent = x
        y.mark = False

    # merge a node with the doubly linked root list
    def merge_with_root_list(self, node):
        if self.root_list is None:
            self.root_list = node
        else:
            node.right = self.root_list.right
            node.left = self.root_list
            self.root_list.right.left = node
            self.root_list.right = node

    # merge a node with the doubly linked child list of a root node
    def merge_with_child_list(self, parent, node):
        if parent.child is None:
            parent.child = node
        else:
            node.right = parent.child.right
            node.left = parent.child
            parent.child.right.left = node
            parent.child.right = node

    # remove a node from the doubly linked root list
    def remove_from_root_list(self, node):
        if node == self.root_list:
            self.root_list = node.right
        node.left.right = node.right
        node.right.left = node.left

    # remove a node from the doubly linked child list
    def remove_from_child_list(self, parent, node):
        if parent.child == parent.child.right:
            parent.child = None
        elif parent.child == node:
            parent.child = node.right
            node.right.parent = parent
        node.left.right = node.right
        node.right.left = node.left

def dijkstra(adjList, source, sink = None):
    n = len(adjList)    #intentionally 1 more than the number of vertices, keep the 0th entry free for convenience
    visited = [False]*n
    distance = [float('inf')]*n

    heapNodes = [None]*n
    heap = FibonacciHeap()
    for i in range(1, n):
        heapNodes[i] = heap.insert(float('inf'), i)     # distance, label

    distance[source] = 0
    heap.decrease_key(heapNodes[source], 0)

    while heap.total_nodes:
        current = heap.extract_min().value
        visited[current] = True

        #early exit
        if sink and current == sink:
            break

        for (neighbor, cost) in adjList[current]:
            if not visited[neighbor]:
                if distance[current] + cost < distance[neighbor]:
                    distance[neighbor] = distance[current] + cost
                    heap.decrease_key(heapNodes[neighbor], distance[neighbor])


    return distance

noOfCities = int (input())
numbers = list(map(int, input().split()))
adjList=[[] for t in range(0,noOfCities+1)]
for i in range(1,noOfCities):
    information = list(map(int, input().split()))
    adjList[information[0]].append(tuple([information[1],information[2]]))
    adjList[information[1]].append(tuple([information[0],information[2]]))

totalshortest = 0
for t in range(1,noOfCities+1):
	shortest = math.inf
	answer = dijkstra(adjList,t)
	for i in answer:
		if ((numbers[answer.index(i)-1]==0) and i<shortest):
			shortest = i
	totalshortest=totalshortest+shortest

print (totalshortest)

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

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