結果

問題 No.439 チワワのなる木
ユーザー H3PO4
提出日時 2021-07-10 09:35:47
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 652 ms / 5,000 ms
コード長 984 bytes
コンパイル時間 105 ms
コンパイル使用メモリ 12,928 KB
実行使用メモリ 37,632 KB
最終ジャッジ日時 2024-07-02 01:45:44
合計ジャッジ時間 6,693 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 28
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
from collections import deque

input = sys.stdin.readline

N = int(input())
S = input()
T = [[] for _ in range(N)]
for _ in range(N - 1):
    a, b = (int(x) - 1 for x in input().split())
    T[a].append(b)
    T[b].append(a)

# dfs
dfs_order = []
parent = [None] * N
d = deque([0])
nonvisited = [True] * N
nonvisited[0] = False
while d:
    v = d.pop()
    dfs_order.append(v)
    for x in T[v]:
        if nonvisited[x]:
            d.append(x)
            parent[x] = v
            nonvisited[x] = False

Csum, Wsum = S.count('c'), S.count('w')

# 子から親へ
C, W = [0] * N, [0] * N
ans = 0
for v in reversed(dfs_order):
    if S[v] == 'c':
        C[v] += 1
    else:  # S[v] == 'w'
        W[v] += 1
    for child in T[v]:
        if child == parent[v]:
            continue
        C[v] += C[child]
        W[v] += W[child]
        if S[v] == 'w':
            ans += W[child] * (Csum - C[child])
    if S[v] == 'w':
        ans += C[v] * (Wsum - W[v])
print(ans)
0