Tietoverkko

import collections
# import numpy as np

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

###outputs 88
testInput="""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
# # exampleInput=testInput
# ei = iter(exampleInput.splitlines())
# print(time.perf_counter())
koneet=GetInput(input)
total=0
# print(time.perf_counter())

defa = collections.defaultdict(int)
def combineFork(haarat):
    a=defa.copy()
    for h in haarat:
        for k,v in h[0].items():
            a[k]+=v
    return a
def clampSelfValue(h):
    addTotal=0
    # s=tuple(h.keys())
    # for i in range(len(h)):
    #     for j in range(i):
    #         addTotal+=min(s[j],s[i])*h[s[j]]*h[s[i]]
    s=sorted(h.keys())
    ta=0
    tb=0
    for i in range(len(h)):
        for j in range(i):
            ta+=s[j]*h[s[j]]
        addTotal+=ta*h[s[i]]
    return addTotal
def clampTo(h,sourceSpeed):
    """in pseudocode:
    y = h[0:sourceSpeed]
    y[sourceSpeed]=sum(h[sourceSpeed:]) + 1"""
    y = {sourceSpeed:1}
    for k,v in h.items():
        if sourceSpeed<=k:
            y[sourceSpeed]+=v
        else:
            y[k]=v
    return y

def CombineHaarat(haarat,sourceSpeed):
    largest=sourceSpeed
    addTotal=0
    yhtyhaarat=combineFork(haarat)
    addTotal+=clampSelfValue(yhtyhaarat)
    for h in haarat:
        addTotal-=clampSelfValue(h[0])
        if largest<h[1]:
            largest=h[1]
    if largest<=sourceSpeed:
        y=yhtyhaarat
        y[sourceSpeed]=y.get(sourceSpeed,0)+1
    else:
        y=clampTo(yhtyhaarat,sourceSpeed)
    for k,v in y.items():
        addTotal+=k*v
    return y,addTotal,largest

def Crossroads(kone,source,sourceSpeed):
    haarat=[]
    for k,v in koneet[kone]:
        if k==source:
            continue
        c=Crossroads(k,kone,v)
        haarat.append(c)
    co=CombineHaarat(haarat,sourceSpeed)
    global total
    total+=co[1]
    return co[0],co[2]

Crossroads(1,0,0)
# print(time.perf_counter())
print(total)
# print("980176904750134603 matches?")

Test 1

Group: 1, 2, 3

Verdict: WRONG ANSWER

Test 2

Group: 2, 3

Verdict: WRONG ANSWER

Test 3

Group: 3

Verdict: TIME LIMIT EXCEEDED