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

Solution (Python):
import sys
input = sys.stdin.readline

n = int(input())
color = list(map(lambda x: int(x)-1,input().split()))
out = [0] * n

total = n * ( n + 1)
total //= 2

stacks = [[[0,0]] for i in range(n)]
segs = [[] for i in range(n)]

adj = [[] for i in range(n)]
visited = [False] * n
dfs = []

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

visited[0] = True
dfs.append(0)
stacks[color[0]].append([1,0])
curr = 1

while dfs:
    node = dfs[-1]
    prev, rem = stacks[color[node]][-1]
    diff = (curr - prev - rem)
    out[color[node]] += (diff * (diff + 1))//2
    stacks[color[node]][-1][1] += (diff) 
    
    if adj[node]:
        nex = adj[node].pop()
        if not visited[nex]:
            visited[nex] = True
            dfs.append(nex)
            curr += 1
            stacks[color[nex]].append([curr,0])
    else:
        dfs.pop()
        prev, rem = stacks[color[node]].pop()
        stacks[color[node]][-1][1] += rem + 1
for i in range(n):
    diff = (curr - stacks[i][0][1])
    out[i] += (diff * (diff + 1))//2
    out[i] = total - out[i]

print('\n'.join(map(str,out)))