結果

問題 No.439 チワワのなる木
ユーザー maspy
提出日時 2020-03-07 03:13:46
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 533 ms / 5,000 ms
コード長 1,730 bytes
コンパイル時間 136 ms
コンパイル使用メモリ 12,928 KB
実行使用メモリ 48,996 KB
最終ジャッジ日時 2024-10-14 11:07:21
合計ジャッジ時間 5,784 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 28
権限があれば一括ダウンロードができます

ソースコード

diff # 

#!/usr/bin/env python3
# %%
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines


# %%
N = int(readline())
S = ' ' + readline().decode().rstrip()
m = map(int, read().split())
AB = zip(m, m)

# %%
graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)

# %%
root = 1
parent = [0] * (N + 1)
order = []
stack = [root]
while stack:
    x = stack.pop()
    order.append(x)
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)

# %%
dp1_w = [0] * (N + 1)
dp1_ww = [0] * (N + 1)
for v in order[::-1]:
    p = parent[v]
    dp1_w[p] += dp1_w[v]
    dp1_ww[p] += dp1_ww[v]
    if S[v] == 'w':
        dp1_w[p] += 1
        dp1_ww[p] += dp1_w[v]

dp1_w, dp1_ww

# %%
dp2_w = [0] * (N + 1)
dp2_ww = [0] * (N + 1)
for v in order[1:]:
    p = parent[v]
    if S[p] == 'w' and S[v] == 'w':
        dp2_w[v] = dp2_w[p] + dp1_w[p] - (1 + dp1_w[v]) + 1
        dp2_ww[v] = (dp2_ww[p] + dp2_w[p]) + (dp1_ww[p] + dp1_w[p]) - (dp1_ww[v] + 2 * dp1_w[v] + 1)
    elif S[p] == 'w' and S[v] == 'c':
        dp2_w[v] = dp2_w[p] + dp1_w[p] - (dp1_w[v]) + 1
        dp2_ww[v] = (dp2_ww[p] + dp2_w[p]) + (dp1_ww[p] + dp1_w[p]) - dp1_w[v] - dp1_ww[v]
    elif S[p] == 'c' and S[v] == 'w':
        dp2_w[v] = dp2_w[p] + dp1_w[p] - (1 + dp1_w[v])
        dp2_ww[v] = dp2_ww[p] + dp1_ww[p] - (dp1_ww[v] + dp1_w[v])
    elif S[p] == 'c' and S[v] == 'c':
        dp2_w[v] = dp2_w[p] + dp1_w[p] - dp1_w[v]
        dp2_ww[v] = dp2_ww[p] + dp1_ww[p] - dp1_ww[v]

dp2_w, dp2_ww


# %%
ww = [x + y for x, y in zip(dp1_ww, dp2_ww)]
answer = sum(x * (ch == 'c') for x, ch in zip(ww, S))
print(answer)
0