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:
import sys

sys.setrecursionlimit(10**7)

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

children = []
size = []
path = []

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

	global children
	children = [[] for i in range(n+1)]
	for i in range(n-1):
		children[a[i]].append(b[i])
		children[b[i]].append(a[i])
	global size
	size = [1] * (n+1)
	global path
	path = [0] * (n+1)

def cut_nodes(node, parent, k):
	nodes = []
	if c[node-1] == k:
		nodes.append(node)
		return nodes
	for child in children[node]:
		if child != parent:
			nodes.extend(cut_nodes(child, node, k))
	return nodes

def dfs(node, parent):
	children_size = []

	k = c[node-1]
	for child in children[node]:
		if child != parent:
			dfs(child, node)
			size[node] += size[child]

			children_size.append(size[child])
			nodes = cut_nodes(child, node, k)
			for leaf in nodes:
				children_size[-1] -= size[leaf]
				path[node] += path[leaf]

	path[node] += n - size[node] + 1
	for s in children_size:
		path[node] += (n - size[node]) * s + s
	for i in range(len(children_size)):
		for j in range(i):
			path[node] += children_size[i] * children_size[j]


def get_answer():
	dfs(1, 0)
	for k in range(1, n+1):
		answer = 0
		nodes = cut_nodes(1, 0, k)
		for leaf in nodes:
			answer += path[leaf]
		print(answer)

if __name__ == '__main__':
	format_input()

	get_answer()

Submission status: Time Limit Exceeded

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

n = int(input())
c = [int(ci) - 1 for ci in input().split()]

def iig(vn, en):
    res = [[] for _ in range(vn)]
    for _ in range(en):
        ai, bi = map(lambda x: int(x) - 1, input().split())
        res[ai].append(bi)
        res[bi].append(ai)
    return res

def sep_pc(graph, root=0):
    parent = [-1 for _ in range(n)]
    children = [None for _ in range(n)]
    def f(x, p):
        parent[x] = p
        children[x] = [y for y in graph[x] if y != p]
        for y in graph[x]:
            if y != p:
                f(y, x)
        return
    f(root, -1)
    return parent, children

def itree(n, root=0):
    return sep_pc(iig(n, n - 1), root)

parent, children = itree(n)

def get_tsize(children, root=0):
    res = [0 for _ in children]
    def f(x):
        res[x] = sum((f(y) for y in children[x]), 1)
        return res[x]
    f(root)
    return res

tsize = get_tsize(children)
head = [n for _ in range(n)]
location = [-1 for _ in range(n)]
def fill_seg_size(x):
    color = c[x]
    if location[color] < 0:
        head[color] -= tsize[x]
    else:
        tsize[location[color]] -= tsize[x]
    prev_location = location[color]
    for y in children[x]:
        location[color] = y
        fill_seg_size(y)
    location[color] = prev_location
    return

fill_seg_size(0)

def comb2(a):
    return a * (a + 1) // 2

vs_by_color = [[] for _ in range(n)]
for i in range(n):
    vs_by_color[c[i]].append(i)
tot = comb2(n)


for color in range(n):
    xs = vs_by_color[color]
    ans = tot - comb2(head[color]) - sum(comb2(tsize[y]) for x in vs_by_color[color] for y in children[x])
    print(ans)

Submission status: Time Limit Exceeded

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

n = int(input())
c = [int(ci) - 1 for ci in input().split()]
v = [[] for _ in range(n)]
for _ in range(n - 1):
    ai, bi = map(lambda x: int(x) - 1, input().split())
    v[ai].append(bi)
    v[bi].append(ai)

par = [-1] * n
ch = [[] for _ in range(n)]
def sep_parnch(x=0, parent=-1):
    par[x] = parent
    for i in v[x]:
        if i == parent:
            continue
        ch[x].append(i)
        sep_parnch(i, x)
    return

sep_parnch()

tsize = [0] * n
def fill_tsize(x=0):
    tsize[x] = sum(fill_tsize(y) for y in ch[x]) + 1
    return tsize[x]

fill_tsize()

seg_size = [[0] * len(ch[x]) for x in range(n)]
def fill_seg_size(x=0, clist=[(None, -1)]*n):
    al, aind = clist[c[x]]
    if al is not None:
        al[aind] += tsize[x]
    for i, y in enumerate(ch[x]):
        clist[c[x]] = seg_size[x], i
        fill_seg_size(y, clist)
    clist[c[x]] = al, aind
    return

head_seg = [([0], 0) for _ in range(n)]
fill_seg_size(0, head_seg)

def comb2(a):
    return a * (a - 1) // 2 + a

vcat = [[] for _ in range(n)]
for i in range(n):
    vcat[c[i]].append(i)

tot = comb2(n)
for col, vs in enumerate(vcat):
    ans = sum(comb2(tsize[y] - seg_size[x][i]) for x in vs for i, y in enumerate(ch[x])) + comb2(n - head_seg[col][0][0])
    print(tot - ans)

Submission status: Time Limit Exceeded

Solution code:
import sys

sys.setrecursionlimit(10**7)

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

children = []

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

	global children
	children = [[] for i in range(n+1)]
	for i in range(n-1):
		children[a[i]].append(b[i])
		children[b[i]].append(a[i])

def dfs(node, parent):
	size = [1] * (n+1)
	path = [0] * (n+1)
	children_size = []

	for child in children[node]:
		if child != parent:
			s, p = dfs(child, node)
			for k in range(n+1):
				path[k] += p[k]
				size[k] += s[k]
				if k == c[node-1]:
					children_size.append(s[k])

	k = c[node-1]
	size[k] = 0
	path[k] += n - size[0] + 1
	for s in children_size:
		path[k] += (n - size[0]) * s + s
	for i in range(len(children_size)):
		for j in range(i):
			path[k] += children_size[i] * children_size[j]

	return size, path

def get_answer():
	size, path = dfs(1, 0)

	for k in range(1, n+1):
		print(path[k])

if __name__ == '__main__':
	format_input()

	get_answer()

Submission status: Time Limit Exceeded

Solution code:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np


def solve(N, C, AB):
    def edges_to_array(A, B, N, M):
        head = np.full(N + 1, -1, np.int32)
        nxt = np.empty(M + M, np.int32)
        to = np.empty(M + M, np.int32)
        for i in range(M):
            a = A[i]
            b = B[i]
            nxt[i << 1] = head[a]
            to[i << 1] = b
            head[a] = i << 1
            nxt[i << 1 | 1] = head[b]
            to[i << 1 | 1] = a
            head[b] = i << 1 | 1
        return head, nxt, to

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = np.zeros_like(head)
        ind_L = np.empty_like(head)
        ind_R = np.empty_like(head)
        tour = np.empty(N + N, np.int32)
        stack = np.empty(N + N, np.int32)
        stack[0] = -root
        stack[1] = root
        t = -1
        s = 2
        while s > 0:
            s -= 1
            v = stack[s]
            t += 1
            tour[t] = v
            if v > 0:
                ind_L[v] = t
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack[s] = -w
                    s += 1
                    stack[s] = w
                    s += 1
                    k = nxt[k]
            else:
                ind_R[-v] = t
        return tour, ind_L, ind_R, parent

    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', '(i8,i4[:],i4[:])')(solve)
    cc.compile()


if sys.argv[-1] == 'ONLINE_JUDGE':
    cc_export()
    exit()

N = int(readline())
C = np.zeros(N + 1, np.int32)
C[1:] = np.array(readline().split(), np.int32)
AB = np.array(read().split(), np.int32)
answer = solve(N, C, AB)
print('\n'.join(answer[1:].astype(str)))

Submission status: Time Limit Exceeded

Solution code:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np


def solve(N, C, AB):
    def edges_to_array(A, B, N, M):
        head = np.full(N + 1, -1, np.int32)
        nxt = np.empty(M + M, np.int32)
        to = np.empty(M + M, np.int32)
        for i in range(M):
            a = A[i]
            b = B[i]
            nxt[i << 1] = head[a]
            to[i << 1] = b
            head[a] = i << 1
            nxt[i << 1 | 1] = head[b]
            to[i << 1 | 1] = a
            head[b] = i << 1 | 1
        return head, nxt, to

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = np.zeros_like(head)
        ind_L = np.empty_like(head)
        ind_R = np.empty_like(head)
        tour = np.empty(N + N, np.int32)
        stack = np.empty(N + N, np.int32)
        stack[0] = -root
        stack[1] = root
        t = -1
        s = 2
        while s > 0:
            s -= 1
            v = stack[s]
            t += 1
            tour[t] = v
            if v > 0:
                ind_L[v] = t
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack[s] = -w
                    s += 1
                    stack[s] = w
                    s += 1
                    k = nxt[k]
            else:
                ind_R[-v] = t
        return tour, ind_L, ind_R, parent

    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', '(i8,i4[:],i4[:])')
    cc.compile()


try:
  from my_module import solve
except:
  cc_export()
  exit()
N = int(readline())
C = np.zeros(N + 1, np.int32)
C[1:] = np.array(readline().split(), np.int32)
AB = np.array(read().split(), np.int32)
answer = solve(N, C, AB)
print('\n'.join(answer[1:].astype(str)))

Submission status: Time Limit Exceeded

Solution code:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np


def solve(N, C, AB):
    def edges_to_array(A, B, N, M):
        head = np.full(N + 1, -1, np.int32)
        nxt = np.empty(M + M, np.int32)
        to = np.empty(M + M, np.int32)
        for i in range(M):
            a = A[i]
            b = B[i]
            nxt[i << 1] = head[a]
            to[i << 1] = b
            head[a] = i << 1
            nxt[i << 1 | 1] = head[b]
            to[i << 1 | 1] = a
            head[b] = i << 1 | 1
        return head, nxt, to

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = np.zeros_like(head)
        ind_L = np.empty_like(head)
        ind_R = np.empty_like(head)
        tour = np.empty(N + N, np.int32)
        stack = np.empty(N + N, np.int32)
        stack[0] = -root
        stack[1] = root
        t = -1
        s = 2
        while s > 0:
            s -= 1
            v = stack[s]
            t += 1
            tour[t] = v
            if v > 0:
                ind_L[v] = t
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack[s] = -w
                    s += 1
                    stack[s] = w
                    s += 1
                    k = nxt[k]
            else:
                ind_R[-v] = t
        return tour, ind_L, ind_R, parent

    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', '(i8,i4[:],i4[:])')
    cc.compile()


if sys.argv[-1] == 'ONLINE_JUDGE':
    cc_export()
    exit()

#from my_module import solve
N = int(readline())
C = np.zeros(N + 1, np.int32)
C[1:] = np.array(readline().split(), np.int32)
AB = np.array(read().split(), np.int32)
answer = solve(N, C, AB)
print('\n'.join(answer[1:].astype(str)))

Submission status: Time Limit Exceeded

Solution code:

import sys
input = lambda : sys.stdin.readline().rstrip()

# 小さすぎても怒られるので
from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, 2*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)

def dfs(ns, prev, now, toru, results):
    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 = results[u]  # 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)
    results[now] = end, toru, s

toru = [0] * (n+1)
q = [(1, -1)] # now, prev
args = []
while q:
    u, prev = q.pop()
    args.append((u,prev))
    for v in ns[u]:
        if prev==v:
            continue
        q.append((v, u))
results = {}
for arg in args[::-1]:
    dfs(ns, arg[1], arg[0], toru, results)

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

Submission status: Time Limit Exceeded

Solution code:

import sys
# 小さすぎても怒られるので
from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, 2*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)

def dfs(ns, prev, now, toru, results):
    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 = results[u]  # 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)
    results[now] = end, toru, s

toru = [0] * (n+1)
q = [(1, -1)] # now, prev
args = []
while q:
    u, prev = q.pop()
    args.append((u,prev))
    for v in ns[u]:
        if prev==v:
            continue
        q.append((v, u))
results = {}
for arg in args[::-1]:
    dfs(ns, arg[1], arg[0], toru, results)

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

Submission status: Time Limit Exceeded

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
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)
    sys.stdout.write(str(item) + "\n")

Submission status: Time Limit Exceeded

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
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: Time Limit Exceeded

Solution code:
from functools import lru_cache
import sys

sys.setrecursionlimit(10**7)

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

children = []
parents = []
tree_size = []

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

	global children
	children = [[] for i in range(n+1)]
	for i in range(n-1):
		children[a[i]].append(b[i])
		children[b[i]].append(a[i])
	global parents
	parents = [0] * (n+1)
	global tree_size
	tree_size = [[] for i in range(n+1)]

@lru_cache(maxsize=None)
def dfs(node, parent):
	size = [1] * (n+1)
	for child in children[node]:
		if child != parent:
			s = dfs(child, node)
			for k in range(n+1):
				if c[node-1] == k:
					tree_size[k].append(s[k])
				size[k] += s[k]

	for k in range(n+1):
		if c[node-1] == k:
			size[k] = 0
	global parents
	parents[node] = parent
	return size

def paths(tree_size):
	return (tree_size ** 2 + tree_size) // 2

def get_answer():
	size = dfs(1, 0)
	whole_paths = paths(size[0])

	for k in range(1, n+1):
		tree_size[k].append(size[k])
		complement = 0
		for i in tree_size[k]:
			complement += paths(i)
		answer = whole_paths - complement
		print(answer)

if __name__ == '__main__':
	format_input()

	get_answer()

Submission status: Time Limit Exceeded

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

Solution code:
from collections import Counter
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

N = int(readline())
C = tuple(map(lambda x: int(x)-1, readline().split()))

temp = map(int,read().split())
temp = zip(temp,temp)

adj = [list() for _ in range(N)]
for a,b in temp:
    adj[a-1].append(b-1)
    adj[b-1].append(a-1)


parent = [None]*N
parent[0] = -1

stack = [0]
order = []
while stack:
    p = stack.pop()
    order.append(p)
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            stack.append(c)

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [N*(N+1)//2]*N

for c in reversed(order):
    p = parent[c]
    if p >= 0:
        dp[c][C[c]] = subtreesize[c]+1
        x = subtreesize[c]-dp[c][C[p]]+1
        cnt[C[p]] -= x*(x+1)//2
        subtreesize[p] += subtreesize[c]+1

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None

dp[0][C[0]] = subtreesize[0]+1
for c in range(N):
    x = subtreesize[0]-dp[0][c]+1
    cnt[c] -= x*(x+1)//2

print(*cnt, sep='\n')

Submission status: Time Limit Exceeded

Solution code:
from collections import Counter

N = int(input())
C = tuple(map(lambda x: int(x)-1,input().split()))

adj = [list() 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)


parent = [None]*N
parent[0] = -1

stack = [0]
order = []
while stack:
    p = stack.pop()
    order.append(p)
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            stack.append(c)

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [N*(N+1)//2]*N

for c in reversed(order):
    p = parent[c]
    if p >= 0:
        dp[c][C[c]] = subtreesize[c]+1
        x = subtreesize[c]-dp[c][C[p]]+1
        cnt[C[p]] -= x*(x+1)//2
        subtreesize[p] += subtreesize[c]+1

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None

dp[0][C[0]] = subtreesize[0]+1
for c in range(N):
    x = subtreesize[0]-dp[0][c]+1
    cnt[c] -= x*(x+1)//2

print(*cnt, sep='\n')

Submission status: Time Limit Exceeded

Solution code:
from collections import Counter

N = int(input())
C = tuple(map(lambda x: int(x)-1,input().split()))

adj = [list() 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)

ROOT = N//2

parent = [None]*N
parent[ROOT] = -1

degree = [0]*N
stack = [ROOT]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    p = parent[c]


    if p >= 0:
        dp[c][C[c]] = subtreesize[c]+1
        x = subtreesize[c]-dp[c][C[p]]+1
        cnt[C[p]] += x*(x+1)//2
        subtreesize[p] += subtreesize[c]+1

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

dp[ROOT][C[ROOT]] = subtreesize[ROOT]+1
for c in range(N):
    x = subtreesize[ROOT]-dp[ROOT][c]+1
    cnt[c] += x*(x+1)//2

total = (subtreesize[ROOT]+1)*(subtreesize[ROOT]+2)//2

print(*(total-v for v in cnt), sep='\n')

Submission status: Time Limit Exceeded

Solution code:
from collections import Counter

N = int(input())
C = tuple(map(lambda x: int(x)-1,input().split()))

adj = [list() 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)


parent = [None]*N
parent[0] = -1

degree = [0]*N
stack = [0]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    p = parent[c]


    if p >= 0:
        dp[c][C[c]] = subtreesize[c]+1
        x = subtreesize[c]-dp[c][C[p]]+1
        cnt[C[p]] += x*(x+1)//2
        subtreesize[p] += subtreesize[c]+1

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

dp[0][C[0]] = subtreesize[0]+1
for c in range(N):
    x = subtreesize[0]-dp[0][c]+1
    cnt[c] += x*(x+1)//2

total = (subtreesize[0]+1)*(subtreesize[0]+2)//2

print(*(total-v for v in cnt), sep='\n')

Submission status: Time Limit Exceeded

Solution code:
from collections import Counter

N = int(input())
C = tuple(map(lambda x: int(x)-1,input().split()))

adj = [list() 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)


parent = [None]*N
parent[0] = -1

degree = [0]*N
stack = [0]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    p = parent[c]


    if p >= 0:
        dp[c][C[c]] = subtreesize[c]+1
        x = subtreesize[c]-dp[c][C[p]]+1
        cnt[C[p]] += x*(x+1)//2
        subtreesize[p] += subtreesize[c]+1
        dp[p] += dp[c]
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

dp[0][C[0]] = subtreesize[0]+1
for c in range(N):
    x = subtreesize[0]-dp[0][c]+1
    cnt[c] += x*(x+1)//2

total = (subtreesize[0]+1)*(subtreesize[0]+2)//2

print(*(total-v for v in cnt), sep='\n')