Contest Duration: - (local time) (100 minutes) Back to Home

Submission #12159656

Source Code Expand

Copy
```import sys
from collections import defaultdict

C = (0,) + tuple(map(int, readline().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)
depth = [0] * (N + 1)  # 根を深さ 1 として
order = []
stack = [root]
depth[root] = 1
while stack:
x = stack.pop()
order.append(x)
for y in graph[x]:
if y == parent[x]:
continue
parent[y] = x
stack.append(y)
depth[y] = depth[x] + 1

def merge(A, B):
if len(A) < len(B):
A, B = B, A
for key, cnt in B.items():
A[key] += cnt
return A

dp = [defaultdict(int) for _ in range(N + 1)]  # 部分木内で、消去済の頂点数の分布
size = [0] * (N + 1)  # 部分木の大きさ
answer = [N * (N + 1) // 2] * (N + 1)

for v in order[::-1]:
size[v] += 1
p = parent[v]
dp[v][C[v]] = size[v]
if not p:
break
size[p] += size[v]
c = C[p]
x = size[v] - dp[v][c]
answer[c] -= x * (x + 1) // 2
dp[p] = merge(dp[p], dp[v])

A = dp[root]
for c in range(1, N + 1):
x = N - A[c]
answer[c] -= x * (x + 1) // 2

```

#### Submission Info

Submission Time 2020-04-20 12:47:43+0900 F - path pass i maspy Python (3.8.2) 600 1427 Byte AC 1750 ms 164128 KB

#### Judge Result

Set Name Sample All
Score / Max Score 0 / 0 600 / 600
Status
 AC × 5
 AC × 24
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt
All line_01.txt, line_3_01.txt, random_01.txt, random_02.txt, random_03.txt, random_100_01.txt, random_100_02.txt, random_10_01.txt, random_10_02.txt, random_1_01.txt, random_1_02.txt, random_2_01.txt, random_2_02.txt, random_2_03.txt, random_3_01.txt, random_3_02.txt, random_3_03.txt, sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt, star_01.txt, star_3_01.txt
Case Name Status Exec Time Memory
line_01.txt AC 1052 ms 136468 KB
line_3_01.txt AC 843 ms 128216 KB
random_01.txt AC 1736 ms 164128 KB
random_02.txt AC 1750 ms 163432 KB
random_03.txt AC 1741 ms 160800 KB
random_100_01.txt AC 1449 ms 145816 KB
random_100_02.txt AC 1453 ms 146132 KB
random_10_01.txt AC 1394 ms 136196 KB
random_10_02.txt AC 1428 ms 136132 KB
random_1_01.txt AC 1364 ms 134212 KB
random_1_02.txt AC 1329 ms 133516 KB
random_2_01.txt AC 1396 ms 133256 KB
random_2_02.txt AC 1365 ms 133504 KB
random_2_03.txt AC 1353 ms 133644 KB
random_3_01.txt AC 1371 ms 134136 KB
random_3_02.txt AC 1389 ms 134080 KB
random_3_03.txt AC 1350 ms 133768 KB
sample_01.txt AC 19 ms 9416 KB
sample_02.txt AC 24 ms 9188 KB
sample_03.txt AC 19 ms 9548 KB
sample_04.txt AC 21 ms 9416 KB
sample_05.txt AC 19 ms 9552 KB
star_01.txt AC 1016 ms 132848 KB
star_3_01.txt AC 887 ms 126040 KB