Sadonkorjuu

class Kaupunki:
    def __init__(self, indeksi, tyyppi):
        self.indeksi = indeksi
        self.tyyppi = tyyppi
        self.naapurit = []
        self.etaisyydet = []
        self.tiet = {}

class Tie:
    def __init__(self):
        self.pituus = 0
        self.paadyt = []

def summaa(lista):
    summa = 0
    for jasen in lista:
        summa = summa + jasen
    return summa

def onkolistassa(lista, listb):
    for jasen in lista:
        if not jasen in listb:
            return False
    return True

def syote():
    global n
    n = int(input())
    tyypit = [int(tyyppi) for tyyppi in input().split(" ")]
    
    tiet = []
    for i in range(n - 1):
        userinput = [int(luku) for luku in input().split(" ")]
        tie = Tie()
        tie.paadyt = [indeksi - 1 for indeksi in userinput[:2]]
        tie.pituus = userinput[2]
        tiet.append(tie)
    return tyypit, tiet

def sievenna(tyypit, tiet):
    global n
    kaupungeittain = [[] for i in range(n)]
    for indeksi in range(n-1):
        kaupungeittain[tiet[indeksi].paadyt[0]].append([tiet[indeksi].paadyt[1], tiet[indeksi].pituus])
        kaupungeittain[tiet[indeksi].paadyt[1]].append([tiet[indeksi].paadyt[0], tiet[indeksi].pituus])

    kaupungit = []

    for indeksi in range(n):
        if tyypit[indeksi] == 0:
            kaupunki = Kaupunki(indeksi, 0)
            kaupunki.naapurit = [luvut[0] for luvut in kaupungeittain[indeksi]]
            kaupunki.etaisyydet = [luvut[1] for luvut in kaupungeittain[indeksi]]
        else:
            kaupunki = Kaupunki(indeksi, 1)
            kaupunki.naapurit = [luvut[0] for luvut in kaupungeittain[indeksi]]
            kaupunki.etaisyydet = [luvut[1] for luvut in kaupungeittain[indeksi]]
        kaupungit.append(kaupunki)
    
    for kaupunki in kaupungit:
        kaupunki.naapurit = [kaupungit[indeksi] for indeksi in kaupunki.naapurit]
        for indeksi in range(len(kaupunki.naapurit)):
            kaupunki.tiet[kaupunki.naapurit[indeksi]] = kaupunki.etaisyydet[indeksi]

    return kaupungit

def loyda(kaupunki, kaupungit, selatut = []):
    selatut.append(kaupunki)
    reitit = []
    for naapuri in kaupunki.naapurit:
        matka = kaupunki.tiet[naapuri]
        if naapuri.tyyppi == 0:
            reitit.append(matka)
            continue
        if naapuri in selatut: # jos naapuri on selattu jo, turha jatkaa.
            continue
        if onkolistassa(naapuri.naapurit, selatut): # jos naapurin kaikki naapurit ovat selattuja, turha jatkaa.
            continue
        else:
            lisamatka = loyda(naapuri, kaupungit, selatut = selatut)
            if lisamatka != 0:
                matka = matka + lisamatka    
                reitit.append(matka)
            
    if reitit == []:
        return 0
    else:
        return min(reitit)

def main():
    tyypit, tiet = syote()
    kaupungit = sievenna(tyypit, tiet)
    reitit = []
    for kaupunki in kaupungit:
        if kaupunki.tyyppi == 1:
            lyhin = loyda(kaupunki, kaupungit, selatut = [])
            reitit.append(lyhin)
    print(summaa(reitit))

if __name__ == "__main__":
    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: TIME LIMIT EXCEEDED

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