Problem:


Score : 600 points

Problem Statement

We have a tree with N vertices numbered 1 to N. The i-th edge in this tree connects Vertex a_i and b_i.
Additionally, each vertex is painted in a color, and the color of Vertex i is c_i. Here, the color of each vertex is represented by an integer between 1 and N (inclusive). The same integer corresponds to the same color; different integers correspond to different colors.

For each k=1, 2, ..., N, solve the following problem:

Find the number of simple paths that visit a vertex painted in the color k one or more times.

Note: The simple paths from Vertex u to v and from v to u are not distinguished.

Constraints

1 \leq N \leq 2 \times 10^5

1 \leq c_i \leq N

1 \leq a_i,b_i \leq N

The given graph is a tree.

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N
c_1 c_2 ... c_N
a_1 b_1
:
a_{N-1} b_{N-1}

Output

Print the answers for k = 1, 2, ..., N in order, each in its own line.

Sample Input 1

3
1 2 1
1 2
2 3

Sample Output 1

5
4
0

Let P_{i,j} denote the simple path connecting Vertex i and j.

There are 5 simple paths that visit a vertex painted in the color 1 one or more times:

P_{1,1}\,,\,
P_{1,2}\,,\,
P_{1,3}\,,\,
P_{2,3}\,,\,
P_{3,3} 

There are 4 simple paths that visit a vertex painted in the color 2 one or more times:

P_{1,2}\,,\,
P_{1,3}\,,\,
P_{2,2}\,,\,
P_{2,3} 

There are no simple paths that visit a vertex painted in the color 3 one or more times.  

Sample Input 2

1
1

Sample Output 2

1

Sample Input 3

2
1 2
1 2

Sample Output 3

2
2

Sample Input 4

5
1 2 3 4 5
1 2
2 3
3 4
3 5

Sample Output 4

5
8
10
5
5

Sample Input 5

8
2 7 2 5 4 1 7 5
3 1
1 2
2 7
4 5
5 6
6 8
7 8

Sample Output 5

18
15
0
14
23
0
23
0



Sample input:
3
1 2 1
1 2
2 3


Sample output:
5
4
0

Submission status: Runtime Error

Solution code:
import sys
# 小さすぎても怒られるので
from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, n))

cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)

toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)

for item in toru[1:]:
    print(item)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(max(1000, n))
# 小さすぎても怒られるので

from collections import defaultdict
n = int(input())
cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)

toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)

for item in toru[1:]:
    print(item)

Submission status: Runtime Error

Solution code:
from collections import defaultdict
n = int(input())
cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)

toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)

for item in toru[1:]:
    print(item)

Submission status: Runtime Error

Solution code:
import asyncio

N = int(input())
Cs = list(map(int, input().split()))
adj = [[] for _ in range(N)]
for _ in range(N - 1):
    a, b = map(int, input().split())
    adj[a - 1].append(b - 1)
    adj[b - 1].append(a - 1)

degs = [0] * N
clusters = [[] for _ in range(N)]


async def go(curr, fro):
    col = Cs[curr] - 1
    deg = 1
    clust = {}
    for n in adj[curr]:
        if n == fro:
            continue
        ch_deg, ch_clust = await go(n, curr)
        clusters[col].append(ch_deg - ch_clust.get(col, 0))
        if len(ch_clust) > len(clust):
            clust, ch_clust = ch_clust, clust
        for k, v in ch_clust.items():
            clust[k] = clust.get(k, 0) + v
        deg += ch_deg
    clust[col] = deg
    return deg, clust


_, root_clust = asyncio.get_event_loop().run_until_complete(go(0, -1))
for i in range(N):
    clusters[i].append(N - root_clust.get(i, 0))


for clust in clusters:
    tot = (N * N + N) // 2
    for cnt in clust:
        tot -= (cnt * cnt + cnt) // 2
    print(tot)

Submission status: Runtime Error

Solution code:
import sys
from collections import defaultdict

sys.setrecursionlimit(10 ** 5 + 5)


def dfs(v, p, ccc, links, ans):
    ret_colors = defaultdict(int)
    ret_count = 1
    cv = ccc[v]
    for u in links[v]:
        if u == p:
            continue
        sub_colors, sub_count = dfs(u, v, ccc, links, ans)
        cc = sub_count - sub_colors[cv]
        ans[cv] -= cc * (cc + 1) // 2
        ret_count += sub_count

        if len(ret_colors) < len(sub_colors):
            ret_colors, sub_colors = sub_colors, ret_colors

        for c, cnt in sub_colors.items():
            ret_colors[c] += cnt

    ret_colors[cv] = ret_count

    return ret_colors, ret_count


def solve(n, ccc, links):
    if n == 1:
        return [1]

    all_pair = n * (n + 1) // 2
    ans = [all_pair] * (n + 1)

    colors, count = dfs(0, -1, ccc, links, ans)
    assert count == n

    for c in range(1, n + 1):
        cc = n - colors[c]
        ans[c] -= cc * (cc + 1) // 2

    return ans[1:]


n, *cab = map(int, sys.stdin.buffer.read().split())
ccc = cab[:n]
links = [set() for _ in range(n)]

for a, b in zip(cab[n + 0::2], cab[n + 1::2]):
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)

print('\n'.join(map(str, solve(n, ccc, links))))

Submission status: Runtime Error

Solution code:
import sys
from collections import defaultdict


def dfs(v, p, ccc, links, ans):
    ret_colors = defaultdict(int)
    ret_count = 1
    cv = ccc[v]
    for u in links[v]:
        if u == p:
            continue
        sub_colors, sub_count = dfs(u, v, ccc, links, ans)
        cc = sub_count - sub_colors[cv]
        ans[cv] -= cc * (cc + 1) // 2
        ret_count += sub_count

        if len(ret_colors) < len(sub_colors):
            ret_colors, sub_colors = sub_colors, ret_colors

        for c, cnt in sub_colors.items():
            ret_colors[c] += cnt

    ret_colors[cv] = ret_count

    return ret_colors, ret_count


def solve(n, ccc, links):
    if n == 1:
        return [1]

    all_pair = n * (n + 1) // 2
    ans = [all_pair] * (n + 1)

    colors, count = dfs(0, -1, ccc, links, ans)
    assert count == n

    for c in range(1, n + 1):
        cc = n - colors[c]
        ans[c] -= cc * (cc + 1) // 2

    return ans[1:]


n, *cab = map(int, sys.stdin.buffer.read().split())
ccc = cab[:n]
links = [set() for _ in range(n)]

for a, b in zip(cab[n + 0::2], cab[n + 1::2]):
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)

print('\n'.join(map(str, solve(n, ccc, links))))

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(10**6)
stdin = sys.stdin

ns = lambda: stdin.readline().rstrip()
ni = lambda: int(stdin.readline().rstrip())
nm = lambda: map(int, stdin.readline().split())
nl = lambda: list(map(int, stdin.readline().split()))

n = ni()
color = nl()
g = [list() for _ in range(n)]
for _ in range(n-1):
    a, b = nm()
    g[a-1].append(b-1)
    g[b-1].append(a-1)

size = [1]*n
par = [-1]*n
h = [list() for _ in range(n+1)]
idx = [0]*n
def dfs(v, p):
    h[color[v]].append(v)
    for x in g[v]:
        if x == p: continue
        par[x] = v
        size[v] += dfs(x, v)
        h[color[v]].append(v)
    h[color[v]].append(v)
    return size[v]
dfs(0, -1)
# print(size)
for v in range(n):
    if g[v][idx[v]] == par[v]:
        idx[v] += 1
acc = [0]*(n+10)
ans = [n*(n+1)//2]*(n+1)
for col in range(1, n+1):
    # print(h[col])
    acc[-1] = 0
    p = [-1]
    for v in h[col]:
        if p[-1] == v:
            p.pop()
            if idx[v] >= len(g[v]):
                acc[p[-1]] += size[v]
                continue
            q = size[g[v][idx[v]]] - acc[v]
            idx[v] += 1
            if idx[v] < len(g[v]) and g[v][idx[v]] == par[v]:
              idx[v] += 1
            ans[col] -= q*(q+1)//2
            acc[v] = 0
            # print(v, q)
            p.append(v)
        else:
            p.append(v)
    q = n - acc[-1]
    # print(-1, q)
    ans[col] -= q*(q+1)//2
print(*ans[1:], sep='\n')

Submission status: Runtime Error

Solution code:
import sys

stdin = sys.stdin

ns = lambda: stdin.readline().rstrip()
ni = lambda: int(stdin.readline().rstrip())
nm = lambda: map(int, stdin.readline().split())
nl = lambda: list(map(int, stdin.readline().split()))

n = ni()
color = nl()
g = [list() for _ in range(n)]
for _ in range(n-1):
    a, b = nm()
    g[a-1].append(b-1)
    g[b-1].append(a-1)

size = [1]*n
par = [-1]*n
h = [list() for _ in range(n+1)]
idx = [0]*n
def dfs(v, p):
    h[color[v]].append(v)
    for x in g[v]:
        if x == p: continue
        par[x] = v
        size[v] += dfs(x, v)
        h[color[v]].append(v)
    h[color[v]].append(v)
    return size[v]
dfs(0, -1)
# print(size)
for v in range(n):
    if g[v][idx[v]] == par[v]:
        idx[v] += 1
acc = [0]*(n+10)
ans = [n*(n+1)//2]*(n+1)
for col in range(1, n+1):
    # print(h[col])
    acc[-1] = 0
    p = [-1]
    for v in h[col]:
        if p[-1] == v:
            p.pop()
            if idx[v] >= len(g[v]):
                acc[p[-1]] += size[v]
                continue
            q = size[g[v][idx[v]]] - acc[v]
            idx[v] += 1
            if idx[v] < len(g[v]) and g[v][idx[v]] == par[v]:
              idx[v] += 1
            ans[col] -= q*(q+1)//2
            acc[v] = 0
            # print(v, q)
            p.append(v)
        else:
            p.append(v)
    q = n - acc[-1]
    # print(-1, q)
    ans[col] -= q*(q+1)//2
print(*ans[1:], sep='\n')

Submission status: Runtime Error

Solution code:
import collections
import sys

sys.setrecursionlimit(10 ** 6)

N = int(input())
colors = [int(x) - 1 for x in input().split()]
neighbors = collections.defaultdict(list)
for i in range(N - 1):
    a, b = map(int, input().split())
    neighbors[a - 1].append(b - 1)
    neighbors[b - 1].append(a - 1)

subtree_sizes = [0] * N


def calc_sizes(node, parent):
    sz = 1
    for i in neighbors[node]:
        if i != parent:
            sz += calc_sizes(i, node)
    subtree_sizes[node] = sz
    return sz


calc_sizes(0, -1)
components = {c: [N] for c in range(N)}
components[colors[0]] = [0]


def calc_components(node, parent, color_component):
    if parent >= 0 and colors[parent] != colors[node]:
        color_component[colors[parent]] += 1
        components[colors[parent]].append(subtree_sizes[node])
        components[colors[node]][color_component[colors[node]]] -= subtree_sizes[node]
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node, color_component)
    if parent >= 0 and colors[parent] != colors[node]:
        color_component[colors[node]] -= 1


calc_components(0, -1, {c: 0 for c in colors})


def count_paths(n):
    return n * (n + 1) // 2


total_paths = count_paths(subtree_sizes[0])
for c in range(N):
    print(total_paths - sum(count_paths(size) for size in components[c]))

Submission status: Runtime Error

Solution code:
import collections

N = int(input())
colors = [int(x) - 1 for x in input().split()]
neighbors = collections.defaultdict(list)
for i in range(N - 1):
    a, b = map(int, input().split())
    neighbors[a - 1].append(b - 1)
    neighbors[b - 1].append(a - 1)

subtree_sizes = [0] * N


def calc_sizes(node, parent):
    sz = 1
    for i in neighbors[node]:
        if i != parent:
            sz += calc_sizes(i, node)
    subtree_sizes[node] = sz
    return sz


calc_sizes(0, -1)
components = {c: [subtree_sizes[0]] for c in range(N)}


def calc_components(node, parent):
    color = colors[node]
    if parent >= 0 and colors[parent] != color:
        components[colors[parent]].append(subtree_sizes[node])
    components[color][-1] -= subtree_sizes[node]
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node)


calc_components(0, -1)


def count_paths(n):
    return n * (n + 1) // 2


total_paths = count_paths(subtree_sizes[0])
for c in range(N):
    print(total_paths - sum(count_paths(size) for size in components[c]))

Submission status: Runtime Error

Solution code:
import sys
from collections import Counter, defaultdict, deque

# sys.setrecursionlimit(2 * 10 ** 5)

(N,) = [int(x) for x in input().split()]
colors = [int(x) for x in input().split()]
edges = [[int(x) for x in input().split()] for i in range(N - 1)]
graph = defaultdict(list)
for u, v in edges:
    graph[u].append(v)
    graph[v].append(u)
parent = {1: None}
children = defaultdict(list)


def rootTree(node):
    for nbr in graph[node]:
        if nbr not in parent:
            parent[nbr] = node
            children[node].append(nbr)
            rootTree(nbr)


rootTree(1)


colorlessCount = defaultdict(int)
colorlessTime = defaultdict(int)
t = 0
for c in range(1, N + 1):
    colorlessTime[c] = 0


def buildComps(node):
    global t
    t += 1
    currColor = colors[node - 1]
    colorlessTime[currColor] += 1
    for child in children[node]:
        oldCount = t - colorlessTime[currColor]
        colorlessTime[currColor] = t
        buildComps(child)
        s = t - colorlessTime[currColor]
        colorlessTime[currColor] = t - oldCount
        colorlessCount[currColor] += s * (s + 1) // 2


buildComps(1)
# print(colorToComps)
for c in colorlessTime:
    s = t - colorlessTime[c]
    colorlessCount[c] += s * (s + 1) // 2
maxPairs = N * (N + 1) // 2
for c in range(1, N + 1):
    print(maxPairs - colorlessCount[c])

Submission status: Runtime Error

Solution code:
import sys

stdin = sys.stdin

ns = lambda: stdin.readline().rstrip()
ni = lambda: int(stdin.readline().rstrip())
nm = lambda: map(int, stdin.readline().split())
nl = lambda: list(map(int, stdin.readline().split()))

n = ni()
color = nl()
g = [list() for _ in range(n)]
for _ in range(n-1):
    a, b = nm()
    g[a-1].append(b-1)
    g[b-1].append(a-1)

size = [1]*n
h = [list() for _ in range(n+1)]
sst = [0]*n
def dfs(v, p):
    h[color[v]].append(v)
    for x in g[v]:
        if x == p: continue
        size[v] += dfs(x, v)
    h[color[v]].append(v)
    return size[v]
dfs(0, -1)
# print(size)
acc = [0]*(n+10)
ans = [n*(n+1)//2]*(n+1)
for col in range(1, n+1):
    # print(h[col])
    acc[-1] = 0
    p = [-1]
    for v in h[col]:
        if p[-1] == v:
            p.pop()
            q = size[v] - acc[v] - 1
            ans[col] -= q*(q+1)//2
            # print(v, q)
            acc[p[-1]] += size[v]
        else:
            p.append(v)
    q = n - acc[-1]
    # print(-1, q)
    ans[col] -= q*(q+1)//2
print(*ans[1:], sep='\n')

Submission status: Runtime Error

Solution code:
import numpy as np

ag = list(reversed(np.argsort(aa)))
st = list(reversed(sorted(aa)))
n_min = 0
n_max = n-1
score = 0
for k, v, i in zip(range(len(st)), st, ag):
    if abs(i - n_min) > abs(n_max - i):
        score += v*abs(i - n_min)
        n_min += 1
    elif abs(n_max - i) > abs(i - n_min):
        score += v*abs(n_max - i)
        n_max -= 1
    else:
        if n_min == n_max:
            score += v*abs(n_max - i)
        elif abs(ag[k+1] - n_min) > abs(n_max - ag[k+1]):
            score += v*(n_max - i)
            n_max -= 1
        elif abs(n_max - ag[k+1]) > abs(ag[k+1] - n_min):
            score += v*abs(i - n_min)
            n_min += 1
print(score)

Submission status: Runtime Error

Solution code:
def n0():return int(input())
def n1():return [int(x) for x in input().split()]
def n2(n):return [int(input()) for _ in range(n)]
def n3(n):return [[int(x) for x in input().split()] for _ in range(n)]

n=n0()
c=n1()
ab=[(x[0],x[1]) for x in n3(n-1)]

import networkx as nx
import matplotlib.pyplot as plt
G = nx.Graph()

G.add_nodes_from(range(1,n+1))
G.add_edges_from(ab)
# nx.draw_networkx(G)
# plt.show()

path=[[x] for x in range(1,n+1)]
for i in range(1,n+1):
    for j in range(i+1,n+1):
        path.extend(list(nx.all_simple_paths(G, source=i, target=j)))
path_color = [[c[x-1] for x in line] for line in path]
dict_color={i:0 for i in range(1,n+1)}
for x in path_color:
    for i in set(x):
        dict_color[i]+=1
for i in dict_color:
    print(dict_color[i])