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)
totalHalf=0
# print(time.perf_counter())

# defa = collections.defaultdict(int)
# def combineFork(haarat):
#     a=defa.copy()
#     for h in haarat:
#         for k,v in h.items():
#             a[k]+=v
#     return a
def addFork(original,haara):
    for k,v in haara.items():
        original[k]=original.get(k,0) + v
    return original
        
def clampSelfValue(h):
    addTotalHalf=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=sum(h.values())
    for k in s:
        # for j in range(i):
        ta+=k*h[k]
        tb-=h[k]
        addTotalHalf+=(ta+tb*k)*h[k]
    return addTotalHalf
def clampSelfValueDifference(h,sourceSpeed):
    "max(h)<sourceSpeed ---> clampSelfValue(h)-clampSelfValue(h+{sourceSpeed:1}) = sourceSpeed"
    # return sum((k*v for k,v in h.items()))*2
    # addTotal=0
    # for k in h:
    #     addTotal+=-k*h[k]
    # return addTotal*2
    # clampSelfValue(h)-clampSelfValue(h+{sourceSpeed:1})
    # ta=sum((k*v for k,v in h.items()))
    addTotalHalf=sourceSpeed
    return addTotalHalf

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(yhtyhaarat,sourceSpeed):
    if not yhtyhaarat:
        return {sourceSpeed:1},sourceSpeed
    if max(yhtyhaarat.keys())<sourceSpeed:
        "addTotalHalf=clampSelfValueDifference(yhtyhaarat,sourceSpeed)"
        yhtyhaarat[sourceSpeed]=1
        return yhtyhaarat,sourceSpeed
    else:
        y=clampTo(yhtyhaarat,sourceSpeed)
        addTotalHalf=clampSelfValue(yhtyhaarat)-clampSelfValue(y)+sum((k*v for k,v in y.items()))*2
        return y,addTotalHalf

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

def CrossroadsLoop(kone,source,sourceSpeed):
    addTotalHalf=0
    path=[(0,0),(0,0),(kone,0)]
    haaratPath=[{},{},{}]
    # for k,v in koneet[kone]:
    #     if k==source:
    #         continue
    #     haarat.append(Crossroads(k,kone,v))
    run=True
    while len(path)>2:
        while koneet[path[-1][0]]:
            newPath,speed = koneet[path[-1][0]].pop()
            if newPath==path[-2][0]:
                pass
            else:
                path.append((newPath,speed))
                haaratPath.append({})
        
        co=CombineHaarat(haaratPath[-1],path[-1][1])
        addFork(haaratPath[-2],co[0])
        addTotalHalf+=co[1]
        haaratPath.pop()
        path.pop()
    global totalHalf
    totalHalf = addTotalHalf
CrossroadsLoop(1,0,0)
# print(time.perf_counter())
totalHalf=totalHalf//2
print(totalHalf)
# print("inaccuracy:",totalHalf-980176904750134603)

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 2, 3

Verdict: ACCEPTED

Test 3

Group: 3

Verdict: TIME LIMIT EXCEEDED