Tietoverkko

 
 
def GetInput():
    n = int(input())
    out=[(n,)]
    for i in range(n-1):
        a = input()
        out.append(tuple(int(b) for b in a.split()))
    return out
exampleInput="""4
1 2 5
2 3 1
2 4 2
"""
exampleInput="""8
1 3 8
2 1 5
4 2 11
5 2 2
6 5 8
7 4 1
8 5 5
"""
exampleInput=[tuple(int(b) for b in a.split()) for a in exampleInput.splitlines()]
I=exampleInput
I=GetInput()
n=I[0][0]
# yhteys voidaan kuvata virtaavan vain yhteen suuntaan? parin suuremmassa yksilössä on referenssi pienempään, 
# ja kun mennään koneiden listan läpi, mennään pienimmästä suurimpaan.
# voidaan huomata että ei ole mahdollista olla ympyröitä, joten jokaisesta koneesta on vain yksi reitti
 
 
import collections
koneet=collections.defaultdict(lambda:{})
 
for i in range(1,n):
    yhteys = I[i]
    b,a=yhteys[0:2]
    koneet[yhteys[0]][yhteys[1]]=yhteys[2]
    koneet[yhteys[1]][yhteys[0]]=yhteys[2]
 
total=0
 
defa = collections.defaultdict(int)
cou = collections.Counter
 
def Crossroads(kone,source,sourceSpeed):
    # sourceSpeed = koneet[kone][source]
    haarat=[]
    yhteydet=defa.copy()
    yhteydet[sourceSpeed]=1
    
    addTotal=0
    for ko_i,ko_i0 in koneet[kone].items():
        if ko_i == source:
            continue
        c=Crossroads(ko_i,kone,ko_i0)
        for j,multiplier in c.items():
            a=min(j,sourceSpeed)
            yhteydet[a]+=multiplier
            addTotal+=multiplier*j
            
        haarat.append(c)
    # for i in range(len(haarat)):
    #     c=haarat[i]
    #     alahaarat=haarat[0:i]
    #     for j,multiplier in c.items():
    #         addTotal+= sum(sum(l0*min(l,j) for l,l0 in haa.items()) for haa in alahaarat)*multiplier
    yhtyhaara=defa.copy()
    for i in range(len(haarat)):
        h=haarat[i]
        for a in h.keys():
            yhtyhaara[a]+=h[a]
    addTotalHalf=0
    soY= sorted(yhtyhaara.keys())
    # for i in range(len(soY)):
    #     # for j in range(len(soY)):
    #     #     addTotalHalf+=min(soY[j],soY[i])*yhtyhaara[soY[i]]*yhtyhaara[soY[j]]
    #     ta=0
    #     tb=0
    #     for j in range(i):
    #         ta+= soY[j]*yhtyhaara[soY[j]]
    #     for j in range(i,len(soY)):
    #         tb+= yhtyhaara[soY[j]]
    #     addTotalHalf+=ta* yhtyhaara[soY[i]]
    #     addTotalHalf+=tb*soY[i]* yhtyhaara[soY[i]]
    ta=0
    tb=0
    for i in soY:
        tb+= yhtyhaara[i]
    for i in soY:
        ta+= i*yhtyhaara[i]
        tb-= yhtyhaara[i]
        addTotalHalf+=(ta+tb*i)* yhtyhaara[i]
 
    for j in range(len(haarat)):
        ta=0
        tb=0
        h=haarat[j]
        soH=sorted(h.keys())
        for i in soH:
            tb+= h[i]
        for i in soH:
            ta+= i*h[i]
            tb-= h[i]
            addTotalHalf-=(ta+tb*i)*h[i]
        
    global total
    total+=addTotal+(addTotalHalf//2)
    
    # yhteydet = {siirtonopeus:määrä}
    # addTotal = summaan lisättävä arvo joka tulee haaran välisistä yhteyksistä, + haarasta aiemmat addTotal
    return yhteydet
 
an = Crossroads(1,0,0)
 
 
# print(I)
# print(koneet)
# print(an)
 
print(total)#-(10**9)

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 2, 3

Verdict: ACCEPTED

Test 3

Group: 3

Verdict: TIME LIMIT EXCEEDED