結果

問題 No.1002 Twotone
ユーザー sasa8uyauya
提出日時 2025-02-28 01:41:45
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 3,097 bytes
コンパイル時間 527 ms
コンパイル使用メモリ 82,824 KB
実行使用メモリ 270,056 KB
最終ジャッジ日時 2025-02-28 01:42:03
合計ジャッジ時間 18,248 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 2 TLE * 2 -- * 29
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
input=sys.stdin.buffer.readline

n,K=map(int,input().split())
e=[[] for i in range(n)]
for _ in range(n-1):
  a,b,c=map(int,input().split())
  a-=1
  b-=1
  c-=1
  e[a]+=[(b,c)]
  e[b]+=[(a,c)]

def CD(e):
  c=[]
  p=[-1]*n
  d=[0]*n
  v=[0]*n
  u=[0]*n
  sq=[(0,n)]
  o=[]
  for start,W in sq:
    q=[start]
    d[start]=0
    while len(q)>0:
      s=q[-1]
      if v[s]==0:
        v[s]=1
        o+=[s]
        for t,_ in e[s]:
          if v[t]==0 and p[t]==-1:
            q+=[t]
            d[t]=d[s]+1
      else:
        u[s]=1+sum(u[t] for t,_ in e[s] if v[t]==0 and p[t]==-1)
        for t,_ in e[s]:
          f=1
          if p[t]==-1:
            if d[t]>d[s]:
              f&=W-u[t]>=u[t]
            else:
              f&=u[s]>=W-u[s]
          if f:
            nc=s
        v[s]=0
        q.pop()
    p[nc]=len(c)
    c+=[nc]
    for t,_ in e[nc]:
      if p[t]==-1:
        if d[t]>d[nc]:
          sq+=[(t,u[t])]
        else:
          sq+=[(t,W-u[nc])]
    o.clear()
  return c,p

cdc,cdp=CD(e)

a1=0
a2=0
a3=0
a4=0
cq1=[0]*K
cqs=0
cq2=[{} for i in range(K)]
sto=[[] for i in range(n)]
cc=[0]*K
cs=set()
v=[0]*n
p=[[] for i in range(n)]
g=[0]*n
for nowc in cdc[::-1]:
  for start,rc in e[nowc]:
    if cdp[start]>cdp[nowc]:
      cc[rc]+=1
      cs.add(rc)
      q=[(start,rc)]
      while len(q)>0:
        s,sc=q[-1]
        if v[s]==0:
          v[s]=1
          sto[start]+=[s]
          if len(cs)<=2:
            p[s]=list(cs).copy()
        while g[s]<len(e[s]):
          t,tc=e[s][g[s]]
          if v[t]==0 and cdp[t]>cdp[nowc] and len(cs)<=3:
            break
          g[s]+=1
        if g[s]<len(e[s]):
          cc[tc]+=1
          if cc[tc]==1:
            cs.add(tc)
          q+=[(t,tc)]
        else:
          cc[sc]-=1
          if cc[sc]==0:
            cs.discard(sc)
          q.pop()
  for start,_ in e[nowc]:
    if cdp[start]>cdp[nowc]:
      for i in sto[start]:
        if len(p[i])==1:
          x=p[i][0]
          cq1[x]+=1
          cqs+=1
        if len(p[i])==2:
          x1,x2=p[i]
          if x2 not in cq2[x1]:
            cq2[x1][x2]=0
          cq2[x1][x2]+=1
  for start,_ in e[nowc]:
    if cdp[start]>cdp[nowc]:
      for i in sto[start]:
        if len(p[i])==1:
          x=p[i][0]
          cq1[x]-=1
          cqs-=1
        if len(p[i])==2:
          x1,x2=p[i]
          cq2[x1][x2]-=1
      for i in sto[start]:
        if len(p[i])==1:
          x=p[i][0]
          a1+=cqs-cq1[x]
        if len(p[i])==2:
          x1,x2=p[i]
          a2+=cq1[x1]+cq1[x2]
          a3+=cq2[x1][x2]
          a4+=1
      for i in sto[start]:
        if len(p[i])==1:
          x=p[i][0]
          cq1[x]+=1
          cqs+=1
        if len(p[i])==2:
          x1,x2=p[i]
          cq2[x1][x2]+=1
  for start,_ in e[nowc]:
    if cdp[start]>cdp[nowc]:
      for i in sto[start]:
        if len(p[i])==1:
          x=p[i][0]
          cq1[x]-=1
          cqs-=1
        if len(p[i])==2:
          x1,x2=p[i]
          cq2[x1][x2]-=1
        v[i]=0
        p[i]=[]
        g[i]=0
      sto[start].clear()

print(a1//2+a2+a3//2+a4)
0