Tietoverkko

import collections


def GetInput(inp):
    n = int(inp())
    # out=collections.defaultdict(lambda:[])
    out=[[] for i in range(n+1)]
    for i in range(n-1):
        # a = inp()
        # c0,c1,c2=(int(b) for b in a.split())
        c0,c1,c2=map(int,inp().split())
        # out[c0][c1]=c2
        # out[c1][c0]=c2
        # c0,c1=sorted(c0,c1)
        out[c0].append((c1,c2))
        out[c1].append((c0,c2))
    return out
# koneetD=collections.defaultdict(lambda:{})

# for i in range(1,n):
#     yhteys = I[i]
#     koneetD[yhteys[0]][yhteys[1]]=yhteys[2]
#     koneetD[yhteys[1]][yhteys[0]]=yhteys[2]
# 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
# """
import time
# with open('C:\\Users\\Okko Heiniö\\Desktop\\Koulu\\Vuosi 3\\Datatähti\\Eexample.txt','r') as fr:
#     exampleInput=fr.read()
#     pass
# ei = iter(exampleInput.splitlines())
# print(time.process_time())
# I=GetInput(ei.__next__)
# koneet=I
koneet=GetInput(input)
# print(time.process_time())
# print(koneet)
total=0
defa = collections.defaultdict(int)
from functools import reduce
def reducer(accumulator, element):
    for key, value in element.items():
        accumulator[key] = accumulator.get(key, 0) + value
    return accumulator
def dict_sum(tests):
    return reduce(reducer, tests, {})

# diag_haaraamount=defa.copy()
def Crossroads(kone,source,sourceSpeed):
    # diag_haaraamount[len(koneet[kone])-1]+=1
    if len(koneet[kone])<=1:
        return {sourceSpeed:1}
    # sourceSpeed = koneet[kone][source]
    haarat=[]
    
    addTotal=0
    for ko_i,ko_i0 in koneet[kone]:
        if ko_i != source:
            haarat.append(Crossroads(ko_i,kone,ko_i0))
            addTotal+=ko_i0
    yhteydet=defa.copy()
    yhteydet[sourceSpeed]=1
    for c in haarat:
        for j,multiplier in c.items():
            a=min(j,sourceSpeed)
            yhteydet[a]+=multiplier
            addTotal+=multiplier*a
    # for j,multiplier in yhteydet.items():
    #     addTotal+=multiplier*j
    
    if len(haarat)>2:
        addTotalHalf=0
        yhtyhaara=dict_sum(haarat)
        soY= sorted(yhtyhaara.keys())
        ta=0
        tb=sum(yhtyhaara.values())
        # 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)):
            h=haarat[j]
            ta=0
            # tb=0
            soH=sorted(h.keys())
            # for i in soH:
            #     tb+= h[i]
            tb=sum(h.values())
            for i in soH:
                ta+= i*h[i]
                tb-= h[i]
                addTotalHalf-=(ta+tb*i)*h[i]
        addTotal+=addTotalHalf//2
    elif len(haarat)==2:
        h0=haarat[0]
        h1=haarat[1]
        h1k=h1.keys()
        # addTotal+=sum(sum(min(i,j)*h0[i]*h1[j] for j in h1k) for i in h0.keys())
        for i in h0.keys():
            ta=0
            for j in h1k:
                ta+=min(i,j)*h1[j]
            addTotal+=ta*h0[i]
            
    global total
    total+=addTotal
    
    # yhteydet = {siirtonopeus:määrä}
    # addTotal = summaan lisättävä arvo joka tulee haaran välisistä yhteyksistä
    return yhteydet

an = Crossroads(1,0,0)


# print(I)
# print(koneet)
# print(an)
# print(diag_haaraamount)

# print(time.process_time())
print(total)#-(10**9)
# print("980176904750134603 matches?")

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 2, 3

Verdict: ACCEPTED

Test 3

Group: 3

Verdict: TIME LIMIT EXCEEDED