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: Time Limit Exceeded

Solution code:
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)]


def go(curr, fro):
    col = Cs[curr] - 1
    deg = 1
    clust = {}
    for n in adj[curr]:
        if n == fro:
            continue
        ch_deg, ch_clust = yield 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


stack = [[None, None], [go(0, -1), None]]
while len(stack) > 1:
    gen, val = stack[-1]
    try:
        params = gen.send(val)
    except StopIteration as e:
        stack[-2][1] = e.value
        stack.pop()
    else:
        stack.append([go(*params), None])

_, root_clust = stack[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: Time Limit Exceeded

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):
    size = 1
    for i in neighbors[node]:
        if i != parent:
            size += calc_sizes(i, node)
    subtree_sizes[node] = size
    return size


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


def calc_components(node, parent):
    color_stack_pushed = False
    if parent >= 0 and colors[parent] != colors[node]:
        color_parent = color_stack[colors[node]][-1]
        components[colors[node]][color_parent] -= subtree_sizes[node]
        components[colors[parent]][node] = subtree_sizes[node]
        color_stack[colors[parent]].append(node)
        color_stack_pushed = True
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node)
    if color_stack_pushed:
        color_stack[colors[parent]].pop()


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].values()))

Submission status: Time Limit Exceeded

Solution code:
from collections import deque

n = 0
c = []
a = []
b = []

def format_input(filename = None):
	global n
	global c, a, b
	if filename == None:
		n = int(input())
		c = list(map(int, input().split()))
		for i in range(n-1):
			ab = list(map(int, input().split()))
			a.append(ab[0])
			b.append(ab[1])

	elif filename == '__random__':
		from random import randint as rng
		n = rng(1, 2 * 10**5)
		c = [rng(1, n) for i in range(n)]
		a = [rng(1, n) for i in range(n-1)]
		b = [rng(1, n) for i in range(n-1)]
		print(n)
		print(' '.join(list(map(str, c))))
		[print(a[i], b[i]) for i in range(n-1)]

def get_answer():
	child = [[] for i in range(n+1)]
	for i in range(n-1):
		child[a[i]].append(b[i])
		child[b[i]].append(a[i])

	complement = [0] * (n+1)

	for k in range(n+1):
		trees = []
		com_mul = 1
		if c.count(k) == 0:
			complement[k] = (n ** 2 + n) // 2
		for i in [child[j+1] for j in range(n) if c[j] == k]:
			for h in i:
				nodes = 0
				k_leaves = 0
				q = deque()
				visited = [False] * (n+1)
				visited[h] = True
				q.append(h)
				while len(q) > 0:
					parent = q.popleft()
					if c[parent-1] == k:
						k_leaves += 1
					else:
						nodes += 1
						for j in child[parent]:
							if not visited[j]:
								visited[j] = True
								q.append(j)
				trees.append([nodes, k_leaves])
				com_mul *= k_leaves

		for i in trees:
			complement[k] += (i[0] ** 2 + i[0]) // 2 * com_mul // i[1]
		complement[k] //= com_mul

	for k in range(1, n+1):
		answer = complement[0] - complement[k]
		print(answer)


if __name__ == '__main__':
	format_input()

	get_answer()

Submission status: Time Limit Exceeded

Solution code:
from collections import deque

n = 0
c = []
a = []
b = []

def format_input(filename = None):
	global n
	global c, a, b
	if filename == None:
		n = int(input())
		c = list(map(int, input().split()))
		for i in range(n-1):
			ab = list(map(int, input().split()))
			a.append(ab[0])
			b.append(ab[1])

	elif filename == '__random__':
		from random import randint as rng
		n = rng(1, 2 * 10**5)
		c = [rng(1, n) for i in range(n)]
		a = [rng(1, n) for i in range(n-1)]
		b = [rng(1, n) for i in range(n-1)]
		print(n)
		print(' '.join(list(map(str, c))))
		[print(a[i], b[i]) for i in range(n-1)]

def get_answer():
	child = [[] for i in range(n+1)]
	for i in range(n-1):
		child[a[i]].append(b[i])
		child[b[i]].append(a[i])

	complement = [0] * (n+1)

	whole_pass = (n ** 2 + n) // 2
	for k in range(1, n+1):
		for i in range(1, n+1):
			q = deque()
			visited = [False] * (n+1)
			visited[i] = True
			q.append(i)
			while len(q) > 0:
				parent = q.popleft()
				if c[parent-1] != k:
					complement[k] += 1
					for j in child[parent]:
						if not visited[j]:
							visited[j] = True
							q.append(j)
		complement[k] += (n - c.count(k))
		complement[k] //= 2


	for k in range(1, n+1):
		answer = whole_pass - complement[k]
		print(answer)


if __name__ == '__main__':
	format_input()

	get_answer()

Submission status: Time Limit Exceeded

Solution code:
import itertools
import networkx as nx
debug = False
if debug:
  import matplotlib.pyplot as plt

N = int(input())
cs = list(map(int, input().split()))
edges = [list(map(int, input().split())) for i in range(1, N)]

g = nx.Graph()
g.add_nodes_from(range(1, N+1), color=cs)
g.add_edges_from(edges)

if debug:
  nx.draw(g, with_labels=True)
  plt.show()

Np = [0] * N
kc_ind = [[i+1 for i, c in enumerate(cs) if c==k] for k in range(1,N+1)]
if debug: print(kc_ind)
combinations = itertools.combinations_with_replacement(range(1, N+1), 2)
for combi in combinations:
  if debug: print(f"combi={combi}")
  paths = nx.all_simple_paths(g, combi[0], combi[1])
  for path in paths:
    if debug: print(f"path={path}")
    for k in range(1,N+1):
      if set(kc_ind[k-1]) & set(list(path)):
        Np[k-1] += 1

for k in range(1,N+1):
  Np[k-1] += len(kc_ind[k-1])

for i in range(N):
  print(Np[i])

Submission status: Time Limit Exceeded

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: Time Limit Exceeded

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)


# cut tree into components that don't have a certain color
# if curr node is color c, each child that isn't color c forms a c-less component.
# Number of c-less path is size * (size + 1) / 2

colorlessCount = defaultdict(int)
colorlessComps = {}
for c in range(1, N + 1):
    colorlessComps[c] = [[]]


def buildComps(node):
    # Grow every colorless comp except for current color
    currColor = colors[node - 1]
    for c in colorlessComps:
        if c != currColor:
            colorlessComps[c][-1].append(node)
    for child in children[node]:
        colorlessComps[currColor].append([])
        buildComps(child)
        temp = colorlessComps[currColor].pop()
        s = len(temp)
        # print('size of colorless', currColor, ' is size', s, 'containing', temp)
        colorlessCount[currColor] += s * (s + 1) // 2


buildComps(1)
# print(colorToComps)
for c in colorlessComps:
    assert len(colorlessComps[c]) == 1
    temp = colorlessComps[c][0]
    s = len(temp)
    # print('size of root colorless', c, ' is size', s, 'containing', temp)
    colorlessCount[c] += s * (s + 1) // 2
# print(colorlessCount)
maxPairs = N * (N + 1) // 2
for c in range(1, N + 1):
    print(maxPairs - colorlessCount[c])

Submission status: Time Limit Exceeded

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
G = nx.Graph()

G.add_nodes_from(range(1,n+1))
G.add_edges_from(ab)

path=[[x] for x in range(1,n+1)]
for i in range(1,n+1):
    path.extend(list(nx.all_simple_paths(G, source=i, target=[x for x in range(i+1,n+1)])))
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])

Submission status: Time Limit Exceeded

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
G = nx.Graph()

G.add_nodes_from(range(1,n+1))
G.add_edges_from(ab)

dict_color={i:0 for i in range(1,n+1)}
for j in c:
    dict_color[j]+=1

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

for i in dict_color:
    print(dict_color[i])

Submission status: Time Limit Exceeded

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
G = nx.Graph()

G.add_nodes_from(range(1,n+1))
G.add_edges_from(ab)

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])

Submission status: Accepted

Solution code:
import sys
from collections import deque

input = sys.stdin.buffer.readline
sys.setrecursionlimit(2*10**5)

N = int(input())
c = list(map(int,input().split()))
edge = [[] for i in range(N+1)]
for _ in range(N-1):
    a,b = map(int,input().split())
    edge[a].append(b)
    edge[b].append(a)


#合流点の検出
deq = deque([1])
parent = [-1 for i in range(N+1)]
parent[1] = 0
topo = []
while deq:
    v = deq.popleft()
    topo.append(v)
    for nv in edge[v]:
        if parent[nv]==-1:
            parent[nv] = v
            deq.append(nv)

topo = topo[::-1]


#辺の張り直し
vertex = []
size = [1 for i in range(N+1)]
size[0] = N

def euler_tour(v,pv):
    vertex.append(v)
    for nv in edge[v]:
        if nv!=pv:
            euler_tour(nv,v)
            size[v] += size[nv]
    vertex.append(v)
euler_tour(1,0)

color_tree_edge = [{c[i]:[]} for i in range(N)]
color_tree_edge = [{c:[] for c in range(1,N+1)}] + color_tree_edge
for v in range(2,N+1):
    color_tree_edge[v][c[parent[v]-1]] = []

parent = [0 for i in range(N+1)]
visit = [None for i in range(N+1)]

for v in vertex:
    if visit[v] is None:
        visit[v] = {}
        for nc in color_tree_edge[v]:
            color_tree_edge[parent[nc]][nc].append(v)
            visit[v][nc] = parent[nc]
            parent[nc] = v
    else:
        for nc in visit[v]:
            parent[nc] = visit[v][nc]

def bfs(color):
    deq = deque([0])
    res = []
    while deq:
        v = deq.popleft()
        res.append(v)
        for nv in color_tree_edge[v][color]:
            deq.append(nv)

    return res[::-1]

ans = [0]*(N+1)

for color in range(1,N+1):
    vertex = bfs(color)
    res = N*(N+1)//2
    for v in vertex:
        if v==0 or c[v-1]!=color:
            tmp = size[v]
            for nv in color_tree_edge[v][color]:
                tmp -= size[nv]
            res -= tmp*(tmp+1)//2
    ans[color] = res

print("\n".join(map(str,ans[1:])))

Submission status: Accepted

Solution code:
# リストを使わないように書き換え

def numba_compile(numba_config):
    import os, sys
    if sys.argv[-1] == "ONLINE_JUDGE":
        from numba import njit
        from numba.pycc import CC
        cc = CC("my_module")
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature)(func)
            cc.export(func.__name__, signature)(func)
        cc.compile()
        exit()
    elif os.name == "posix":
        exec(f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}")
        for func, _ in numba_config:
            globals()[func.__name__] = vars()[func.__name__]
    else:
        from numba import njit
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature, cache=True)(func)
        print("compiled!", file=sys.stderr)

import numpy as np

def sort_edges(N, edges_):
    # N: 頂点番号の最大値
    M = len(edges_)
    edges = np.empty((M * 2, 2), dtype=np.int64)
    edges[:M] = edges_
    edges[M:] = edges_[:, ::-1]
    order = np.argsort(edges[:, 0])  # O(N) にできなくもない
    edges = edges[order, 1]
    c = np.zeros(N+1, dtype=np.int64)
    c_ = np.bincount(edges_.ravel())  # minlength を使わせて
    c[:len(c_)] = c_
    c = np.cumsum(c)
    lefts = np.zeros(len(c) + 1, dtype=np.int64)
    lefts[1:] = c
    return edges, lefts

def eular_tour(edges, lefts, root):
    # グラフは 1-indexed が良い
    n = len(lefts)-1
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        p = parents[v]
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in edges[lefts[v]:lefts[v+1]]:
            if p != u:
                parents[u] = v
                stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

def solve(N, C, AB):
    edges, lefts = sort_edges(N, AB)
    tour, firsts, lasts, parents = eular_tour(edges, lefts, 1)
    Siz = lasts - firsts + 1 >> 1
    Cnt = np.zeros(N+1, dtype=np.int64)
    Ans = np.zeros(N+1, dtype=np.int64)
    Cnt_old = np.empty_like(Cnt)
    Siz_tmp = np.full_like(Cnt, 1)
    for v in tour:
        if v >= 0:
            p = parents[v]
            if p != -100:
                Cnt[C[p]] = N - Siz[v]  # 上を塞ぐ

            color = C[v]
            Cnt_old[v] = Cnt[color]
        else:
            v = -v
            Ans[C[v]] += N - Siz_tmp[v] + 1 - Cnt_old[v]
            Cnt[C[v]] = Cnt_old[v] + Siz_tmp[v]

            p = parents[v]
            if p != -100:
                Siz_tmp[p] += Siz[v]
                Ans[C[p]] += Siz[v] * (N - Siz_tmp[p] + 1 - Cnt_old[p])

    return Ans[1:]


numba_compile([
    [sort_edges, "Tuple((i8[:],i8[:]))(i8,i8[:,:])"],
    [eular_tour, "Tuple((i8[:],i8[:],i8[:],i8[:]))(i8[:],i8[:],i8)"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])


import sys
def main():
    input = sys.stdin.buffer.readline
    read = sys.stdin.buffer.read
    N = int(input())
    C = np.empty(N+1, dtype=np.int64)
    C[1:] = input().split()
    AB = np.array(read().split(), dtype=np.int64).reshape(N-1, 2)
    Ans = solve(N, C, AB)
    print("\n".join(map(str, Ans.tolist())))

main()

Submission status: Accepted

Solution code:
# 入力のみ
# 入力をいい感じに

def numba_compile(numba_config):
    import os, sys
    if sys.argv[-1] == "ONLINE_JUDGE":
        from numba import njit
        from numba.pycc import CC
        cc = CC("my_module")
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature)(func)
            cc.export(func.__name__, signature)(func)
        cc.compile()
        exit()
    elif os.name == "posix":
        exec(f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}")
        for func, _ in numba_config:
            globals()[func.__name__] = vars()[func.__name__]
    else:
        from numba import njit
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature, cache=True)(func)
        print("compiled!", file=sys.stderr)

import numpy as np

def eular_tour(G, root):
    # G は破壊的に有向木にされる
    # グラフは 1-indexed が良い
    n = len(G)
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            parents[u] = v
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

def solve(N, C, AB):
    G = [[0]*0 for _ in range(N+1)]
    for i in range(N-1):
        a, b = AB[i]
        G[a].append(b)
        G[b].append(a)
    tour, firsts, lasts, parents = eular_tour(G, 1)
    Siz = lasts - firsts + 1 >> 1

    Cnt = np.zeros(N+1, dtype=np.int64)
    Ans = np.zeros(N+1, dtype=np.int64)
    Cnt_old = np.empty_like(Cnt)
    Siz_tmp = np.full_like(Cnt, 1)
    for v in tour:
        if v >= 0:
            p = parents[v]
            if p != -100:
                Cnt[C[p]] = N - Siz[v]  # 上を塞ぐ

            color = C[v]
            Cnt_old[v] = Cnt[color]
        else:
            v = -v
            Ans[C[v]] += N - Siz_tmp[v] + 1 - Cnt_old[v]
            Cnt[C[v]] = Cnt_old[v] + Siz_tmp[v]

            p = parents[v]
            if p != -100:
                Siz_tmp[p] += Siz[v]
                Ans[C[p]] += Siz[v] * (N - Siz_tmp[p] + 1 - Cnt_old[p])

    return Ans[1:]


numba_compile([
    [eular_tour, "Tuple((i8[:],i8[:],i8[:],i8[:]))(List(List(i8)),i8)"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])


import sys
def main():
    input = sys.stdin.buffer.readline
    read = sys.stdin.buffer.read
    N = int(input())
    C = np.empty(N+1, dtype=np.int64)
    C[1:] = input().split()
    AB = np.array(read().split(), dtype=np.int64).reshape(N-1, 2)
    Ans = solve(N, C, AB)
    print("\n".join(map(str, Ans.tolist())))

main()

Submission status: Accepted

Solution code:
# tolist つけ忘れ
# https://atcoder.jp/contests/abc163/submissions/16271142

def numba_compile(numba_config):
    import os, sys
    if sys.argv[-1] == "ONLINE_JUDGE":
        from numba import njit
        from numba.pycc import CC
        cc = CC("my_module")
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature)(func)
            cc.export(func.__name__, signature)(func)
        cc.compile()
        exit()
    elif os.name == "posix":
        exec(f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}")
        for func, _ in numba_config:
            globals()[func.__name__] = vars()[func.__name__]
    else:
        from numba import njit
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature, cache=True)(func)
        print("compiled!", file=sys.stderr)

import numpy as np

def eular_tour(G, root):
    # G は破壊的に有向木にされる
    # グラフは 1-indexed が良い
    n = len(G)
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            parents[u] = v
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

def solve(N, C, AB):
    G = [[0]*0 for _ in range(N+1)]
    for i in range(N-1):
        a, b = AB[i]
        G[a].append(b)
        G[b].append(a)
    tour, firsts, lasts, parents = eular_tour(G, 1)
    Siz = lasts - firsts + 1 >> 1

    Cnt = np.zeros(N+1, dtype=np.int64)
    Ans = np.zeros(N+1, dtype=np.int64)
    Cnt_old = np.empty_like(Cnt)
    Siz_tmp = np.full_like(Cnt, 1)
    for v in tour:
        if v >= 0:
            p = parents[v]
            if p != -100:
                Cnt[C[p]] = N - Siz[v]  # 上を塞ぐ

            color = C[v]
            Cnt_old[v] = Cnt[color]
        else:
            v = -v
            Ans[C[v]] += N - Siz_tmp[v] + 1 - Cnt_old[v]
            Cnt[C[v]] = Cnt_old[v] + Siz_tmp[v]

            p = parents[v]
            if p != -100:
                Siz_tmp[p] += Siz[v]
                Ans[C[p]] += Siz[v] * (N - Siz_tmp[p] + 1 - Cnt_old[p])

    return Ans[1:]


numba_compile([
    [eular_tour, "Tuple((i8[:],i8[:],i8[:],i8[:]))(List(List(i8)),i8)"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])


import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = np.array(["0"] + input().split(), dtype=np.int64)
AB = np.array([list(map(int, input().split())) for _ in range(N-1)], dtype=np.int64)
Ans = solve(N, C, AB)
print("\n".join(map(str, Ans.tolist())))

Submission status: Accepted

Solution code:
# tolist つけ忘れ
# https://atcoder.jp/contests/abc163/submissions/16270790

def numba_compile(numba_config):
    import os, sys
    if sys.argv[-1] == "ONLINE_JUDGE":
        from numba import njit
        from numba.pycc import CC
        cc = CC("my_module")
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature)(func)
            cc.export(func.__name__, signature)(func)
        cc.compile()
        exit()
    elif os.name == "posix":
        exec(f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}")
        for func, _ in numba_config:
            globals()[func.__name__] = vars()[func.__name__]
    else:
        from numba import njit
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature, cache=True)(func)
        print("compiled!", file=sys.stderr)

import numpy as np

def dfs1(v, G, Siz):
    siz = 1
    for u in G[v]:
        G[u].remove(v)
        siz += dfs1(u, G, Siz)
    Siz[v] = siz
    return siz

def dfs(v, N, C, G, Cnt, Siz, Ans):
    color = C[v]
    cnt_old = Cnt[color]
    siz = 1
    siz_up = N
    for u in G[v]:
        Cnt[color] = N - Siz[u]  # 上を塞ぐ
        s = dfs(u, N, C, G, Cnt, Siz, Ans)
        siz += s
        siz_up -= s
        Ans[color] += s * (siz_up - cnt_old)
    Ans[color] += siz_up - cnt_old
    Cnt[color] = cnt_old + siz
    return siz

def solve(N, C, AB):
    G = [[0]*0 for _ in range(N)]
    for i in range(N-1):
        a, b = AB[i]
        G[a].append(b)
        G[b].append(a)
    Cnt = np.zeros(N, dtype=np.int64)
    Ans = np.zeros(N, dtype=np.int64)
    Siz = np.zeros(N, dtype=np.int64)
    dfs1(0, G, Siz)
    dfs(0, N, C, G, Cnt, Siz, Ans)
    return Ans

numba_compile([
    [dfs1, "i8(i8,List(List(i8)),i8[:])"],
    [dfs, "i8(i8,i8,i8[:],List(List(i8)),i8[:],i8[:],i8[:])"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])

import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = np.array(input().split(), dtype=np.int64) - 1
AB = np.array([list(map(int, input().split())) for _ in range(N-1)], dtype=np.int64) - 1
Ans = solve(N, C, AB)
print("\n".join(map(str, Ans.tolist())))

Submission status: Accepted

Solution code:
# 全部非再帰に書き換え

def numba_compile(numba_config):
    import os, sys
    if sys.argv[-1] == "ONLINE_JUDGE":
        from numba import njit
        from numba.pycc import CC
        cc = CC("my_module")
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature)(func)
            cc.export(func.__name__, signature)(func)
        cc.compile()
        exit()
    elif os.name == "posix":
        exec(f"from my_module import {','.join(func.__name__ for func, _ in numba_config)}")
        for func, _ in numba_config:
            globals()[func.__name__] = vars()[func.__name__]
    else:
        from numba import njit
        for func, signature in numba_config:
            globals()[func.__name__] = njit(signature, cache=True)(func)
        print("compiled!", file=sys.stderr)

import numpy as np

def eular_tour(G, root):
    # G は破壊的に有向木にされる
    # グラフは 1-indexed が良い
    n = len(G)
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            parents[u] = v
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

def solve(N, C, AB):
    G = [[0]*0 for _ in range(N+1)]
    for i in range(N-1):
        a, b = AB[i]
        G[a].append(b)
        G[b].append(a)
    tour, firsts, lasts, parents = eular_tour(G, 1)
    Siz = lasts - firsts + 1 >> 1

    Cnt = np.zeros(N+1, dtype=np.int64)
    Ans = np.zeros(N+1, dtype=np.int64)
    Cnt_old = np.empty_like(Cnt)
    Siz_tmp = np.full_like(Cnt, 1)
    for v in tour:
        if v >= 0:
            p = parents[v]
            if p != -100:
                Cnt[C[p]] = N - Siz[v]  # 上を塞ぐ

            color = C[v]
            Cnt_old[v] = Cnt[color]
        else:
            v = -v
            Ans[C[v]] += N - Siz_tmp[v] + 1 - Cnt_old[v]
            Cnt[C[v]] = Cnt_old[v] + Siz_tmp[v]

            p = parents[v]
            if p != -100:
                Siz_tmp[p] += Siz[v]
                Ans[C[p]] += Siz[v] * (N - Siz_tmp[p] + 1 - Cnt_old[p])

    return Ans[1:]


numba_compile([
    [eular_tour, "Tuple((i8[:],i8[:],i8[:],i8[:]))(List(List(i8)),i8)"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])


import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = np.array(["0"] + input().split(), dtype=np.int64)
AB = np.array([list(map(int, input().split())) for _ in range(N-1)], dtype=np.int64)
Ans = solve(N, C, AB)
print("\n".join(map(str, Ans)))