結果

問題 No.1418 Sum of Sum of Subtree Size
ユーザー lam6er
提出日時 2025-03-20 18:51:13
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 249 ms / 2,000 ms
コード長 1,184 bytes
コンパイル時間 251 ms
コンパイル使用メモリ 82,108 KB
実行使用メモリ 102,920 KB
最終ジャッジ日時 2025-03-20 18:52:49
合計ジャッジ時間 7,362 ms
ジャッジサーバーID
(参考情報)
judge2 / judge3
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 41
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
sys.setrecursionlimit(1 << 25)

n = int(sys.stdin.readline())
edges = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b = map(int, sys.stdin.readline().split())
    edges[a].append(b)
    edges[b].append(a)

sum_edges = 0
size = [0] * (n + 1)

# Using iterative post-order traversal to compute sizes and sum_edges
stack = []
stack.append((1, -1, False))  # (node, parent, visited)

while stack:
    node, parent, visited = stack.pop()
    if not visited:
        # Push back the node as visited after processing children
        stack.append((node, parent, True))
        # Add children in reverse order to process them in order
        for child in reversed(edges[node]):
            if child != parent:
                stack.append((child, node, False))
    else:
        # Calculate the size of the current node
        s = 1
        for child in edges[node]:
            if child != parent:
                s += size[child]
        size[node] = s
        # Add contribution for each child edge
        for child in edges[node]:
            if child != parent:
                sum_edges += size[child] * (n - size[child])

ans = 2 * sum_edges + n * n
print(ans)
0