提出 #24947611

---

ソースコード 拡げる

import sys
from collections import defaultdict, deque

def main(f):
  N = int(f.readline())
  link = [[] for _ in range(N+1)]
  for _ in range(N-1):
    a, b = map(int, f.readline().split())
    link[a].append(b)
    link[b].append(a)

# dp[n] = (c, d) : n を頂点とするサブツリーについての (ノード数, 「子孫 i から n の距離」の合計)
  dp = [(0, 0)] * (N+1)
  dists = 0

  q = deque()
  q.append((1, 0)) # (node, parent)
  while q:
    i, p = q.pop()
    if i < 0:
      c, d = 0, 0
      for j in link[-i]:
        if j == p:
          continue
        c_j, d_j = dp[j]
        c += c_j
        d += d_j + 1*(c_j) # j->i の距離 1 を各子孫について加える
      dp[-i] = (c+1, d)

      links = []
      for j in link[-i]:
        if j == p:
          continue
        links.append(j)

      # 部分和 c_(j+1) + c_(j+2) +... を計算するために、事前に部分和 c_1 + c_2 + ... + c_n を計算
      # 部分和 d_(j+1) + d_(j+2) +... を計算するために、事前に部分和 d_1 + d_2 + ... + d_n を計算
      c_sum, d_sum = [0], [0]
      for index in range(len(links)): 
        c_j, d_j = dp[links[index]]
        c_sum.append(c_sum[-1] + c_j)
        d_sum.append(d_sum[-1] + d_j)
        d += d_j
        c += c_j

      for index in range(len(links)):
        c_j, d_j = dp[links[index]]
        dists += (d_j+c_j)*(c_sum[-1]-c_sum[index+1]) + c_j*((d_sum[-1]-d_sum[index+1])+(c_sum[-1]-c_sum[index+1]))
      continue

    q.append((-i, p)) # 帰りがけ順で処理するための番兵

    for j in link[i]:
      if j == p:
        continue
      q.append((j, i))

  for n in range(1, N+1):
    dists += dp[n][1] # k -> i の距離をすべて加える
  print(dists)

main(sys.stdin)

提出情報

ジャッジ結果