Sadonkorjuu

import math
 
 
def I():
    return map(int, input().split())
 
 
a = """6
1 1 0 0 1 1
1 2 20
2 3 30
2 5 50
3 6 60
1 4 10
""".splitlines()
# print(a)
# input = iter(a).__next__
 
 
def main():
    n, = I()
    cities = list(I())
    connections = [dict() for _ in range(n)]
    for i in range(n - 1):
        a, b, c = I()
        connections[a - 1][b - 1] = c
        connections[b - 1][a - 1] = c
    head = Node(None, 0, math.inf, cities[0])
    city = head
    while city is not None:
        city.form_children(cities, connections)
        city, going_up = city.next()
    city = head
    while city is not None:
        new, going_up = city.next()
        if new is None:
            # city is the root, and has finished.
            break
        elif city.marker == 0:  # new is city's parent
            assert going_up
            distance = city.distance
            if new.closest > city.closest + distance:
                new.closest = city.closest + distance
                new.marker = 0  # resetting the node
            city = new
        else:  # city is new's parent
            assert not going_up
            distance = new.distance
            if new.closest > city.closest + distance:
                new.closest = city.closest + distance
                new.checked = False
            if not new.checked:
                city = new
                new.checked = True
 
    city = head
    while city is not None:
        city.confirm()
        city, going_up = city.next()
 
    total = 0
    city = head
    while city is not None:
        new, going_up = city.next()
        if going_up:
            total += city.closest
        city = new
    print(total)
    # print(head.view())
 
 
class Node:
    def __init__(self, parent: "Node|None", id: int, distance: float, city_type: bool):
        self.parent = parent
        # if parent is not None:
        #     parent.children.append(self)
        self.children: list[Node] = []
        self.distance = distance
        self.city_type = city_type
        self.id = id
        self.closest = math.inf if city_type else 0
        self.formed_children = False
        self.checked = False
 
    def form_children(self, cities: list[bool], connections: list[dict[int, float]]):
        if self.formed_children:
            return
        for city_id, distance in connections[self.id].items():
            if (not self.parent) or city_id != self.parent.id:
                self.children.append(Node(self, city_id, distance, cities[city_id]))
        self.marker = 0
        self.formed_children = True
 
    def next(self):
        "returns the next node, and True if going up, False is going down"
        if self.marker < len(self.children):
            self.marker += 1
            return self.children[self.marker - 1], False
        self.marker = 0
        return self.parent, True
 
    def confirm(self):
        assert all(self.closest <= c.closest + c.distance for c in self.children)
        assert self.parent is None or self.closest <= self.parent.closest + self.distance
 
    def view(self, indent: int = 0):
        out = ""
        out += " " * indent + f"{self.id}: {self.closest}, {self.distance}" + "\n"
        for c in self.children:
            out += c.view(indent + 4)
        return out
 
 
main()

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

Test 16

Group: 1, 2

Verdict: ACCEPTED

Test 17

Group: 1, 2

Verdict: ACCEPTED

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