結果

問題 No.1333 Squared Sum
ユーザー persimmon-persimmon
提出日時 2021-02-14 14:14:52
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,310 ms / 2,000 ms
コード長 1,796 bytes
コンパイル時間 282 ms
コンパイル使用メモリ 82,432 KB
実行使用メモリ 205,452 KB
最終ジャッジ日時 2024-07-21 22:38:31
合計ジャッジ時間 34,812 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 44
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
input=sys.stdin.readline
def main2(n,uvw):
  ki=[[] for _ in range(n)]
  for u,v,w in uvw:
    u,v=u-1,v-1
    ki[u].append([v,w])
    ki[v].append([u,w])
  mod=10**9+7
  def dfs_tree(n,root,ki):
    q=[root]
    parent=[None]*n
    dist=[0]*n
    for v in q:
      for nv,cost in ki[v]:
        if nv==root or parent[nv] is not None:continue
        dist[nv]=dist[v]+cost
        parent[nv]=v,cost
        q.append(nv)
    return parent,q,dist

  parent,order,dist=dfs_tree(n,0,ki)
  dp0=[[] for _ in range(n)]
  for v in order[::-1]:
    numv,sumv=1,0
    for num_,sum_,cost_ in dp0[v]:
      numv+=num_
      sumv+=sum_+num_*cost_
      sumv%=mod
    dp0[v]=numv,sumv
    if v==0:break
    p,cost=parent[v]
    dp0[p].append([numv,sumv,cost])
  # dp0[v]:親からvへの矢印に対応する部分木の値
  # dp1[v]:vから親への矢印に対応する部分木の値
  ans=0
  dp0[0]=None
  dp1=[None]*n
  parent[0]=[-1,-1]
  for v in order:
    snumv,ssumv=1,0
    for nv,cost in ki[v]:
      if nv==parent[v][0]:
        num_,sum_=dp1[v]
      else:
        num_,sum_=dp0[nv]
      snumv+=num_
      ssumv+=sum_+num_*cost
      ssumv%=mod
    for nv,cost in ki[v]:
      #if nv==0:continue
      if nv==parent[v][0]:
        num0,sum0=dp1[v]
      else:
        num0,sum0=dp0[nv]
      num1,sum1=snumv-num0,ssumv-sum0-cost*num0
      num1,sum1=num1%mod,sum1%mod
      tmp=num0*sum1+num1*sum0+cost*num0*num1
      tmp*=cost
      tmp%=mod
      ans+=tmp
      ans%=mod
      dp1[nv]=num1,sum1
  #print('dp0')
  #for i in range(n):print(i,dp0[i])
  #print('dp1')
  #for i in range(n):print(i,dp1[i])
  ans*=pow(2,mod-2,mod)
  return ans%mod

if __name__=='__main__':
  n=int(input())
  uvw=[list(map(int,input().split())) for _ in range(n-1)]
  ret2=main2(n,uvw)
  print(ret2)
0