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
sys.setrecursionlimit(10**6)
def main():
    n=int(input())
    color=list(map(int,input().split()))
    nodes=[[]for i in range(n)]
    cut=[[0]for i in range(n+1)]
    ans=[0]*(n+1)
    for i in range(n-1):
        a,b=map(int,input().split())
        nodes[a-1].append(b-1)
        nodes[b-1].append(a-1)
    def num(p,parent):
        s=0
        for child in nodes[p]:
            if child==parent:
                continue
            cut[color[p]].append(0)
            nc=num(child,p)
            group=nc-cut[color[p]].pop()
            ans[color[p]]+=group*(group+1)//2
            s+=nc
        s+=1
        cut[color[p]][-1]+=s
        return s
    num(0,-1)
    for a,c in zip(ans[1:],cut[1:]):
        group=n-c[0]
        c=group*(group+1)//2
        print(n*(n+1)//2-c-a)
main()

Solution (Python):
import sys
sys.setrecursionlimit(10**6)
def main():
    n=int(input())
    color=list(map(int,input().split()))
    nodes=[[]for i in range(n)]
    cut=[[0]for i in range(n+1)]
    ans=[0]*(n+1)
    for i in range(n-1):
        a,b=map(int,input().split())
        nodes[a-1].append(b-1)
        nodes[b-1].append(a-1)
    def num(p,parent):
        s=0
        for child in nodes[p]:
            if child==parent:
                continue
            cut[color[p]].append(0)
            nc=num(child,p)
            group=nc-cut[color[p]].pop()
            ans[color[p]]+=group*(group+1)//2
            s+=nc
        s+=1
        cut[color[p]][-1]+=s
        return s
    num(0,-1)
    for a,c in zip(ans[1:],cut[1:]):
        group=n-c[0]
        c=group*(group+1)//2
        print(n*(n+1)//2-c-a)
main()

Solution (Python):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 6)
from collections import defaultdict

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

ret = [N * (N - 1) // 2 + N] * N

par = [[N + i] for i in range(N)]
vec2 = [[] for _ in range(N * 2)]
have = [0] * (N * 2)
way = [0] * (N * 2)
fr = [0] * (N + 2)

for i in range(N, 2 * N) :
  have[i] = N + 1

mem = defaultdict(int)

def dfs(cur, pre) :
  s = 1
  ancestor = par[c[cur]][-1]
  vec2[ancestor].append(cur)
  if pre != -1 :
    fr[cur] = way[ancestor]
  par[c[cur]].append(cur)
  for nex in vec[cur] :
    if nex == pre : continue
    way[cur] = nex
    s += dfs(nex, cur)
    mem[(cur, nex)] = have[nex]
  par[c[cur]].pop()
  have[cur] = s
  return s
  
dfs(0, -1)

for i in range(N, 2 * N) :
  mem[(i, 0)] = N

for i in range(N) :
  q = [(N + i, -1)]
  while q :
    cur, pre = q.pop()
    for nex in vec2[cur] :
      mem[(cur, fr[nex])] -= have[nex]
      q.append((nex, cur))

for k, v in mem.items() :
  if v <= 0: continue
  if k[0] >= N :
    ret[k[0] % N] -= v * (v - 1) // 2 + v
  else :
    ret[c[k[0]]] -= v * (v - 1) // 2 + v
  
for i in range(N) :
  print(ret[i])

Solution (Python):
mod = 1000000007
eps = 10**-9


def main():
    import sys
    from collections import deque
    input = sys.stdin.readline

    def calc(n):
        return (n * (n-1)) // 2 + n

    def merge(A, B):
        if len(A) > len(B):
            A, B = B, A
        for a in A:
            if a in B:
                B[a] += A[a]
            else:
                B[a] = A[a]
        return B

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

    que = deque()
    que.append(1)
    seen = [-1] * (N+1)
    seen[1] = 0
    par = [0] * (N+1)
    child = [[] for _ in range(N+1)]
    seq = []
    while que:
        v = que.popleft()
        seq.append(v)
        for u in adj[v]:
            if seen[u] == -1:
                seen[u] = seen[v] + 1
                par[u] = v
                child[v].append(u)
                que.append(u)
    seq.reverse()

    size = [1] * (N+1)
    dic = [{} for _ in range(N+1)]
    ans = [calc(N)] * (N+1)
    for v in seq:
        for u in child[v]:
            size[v] += size[u]
        c = C[v-1]
        for u in child[v]:
            if c in dic[u]:
                ans[c] -= calc(size[u] - dic[u][c])
            else:
                ans[c] -= calc(size[u])
            dic[v] = merge(dic[v], dic[u])
        dic[v][c] = size[v]
    for c in range(1, N+1):
        if c in dic[1]:
            print(ans[c] - calc(N - dic[1][c]))
        else:
            print(ans[c] - calc(N))


if __name__ == '__main__':
    main()

Solution (Python):
import sys
from collections import defaultdict


sys.setrecursionlimit(10**6)

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


def dfs(v,p):
    ret=defaultdict(int)
    size=1
    for vv in g[v]:
        if vv==p:
            continue
        ss,d=dfs(vv,v)
        size+=ss
        ans[c[v]]+=sum(ss-d[c[v]])
        
        
        if len(ret)<len(d):
            ret,d=d,ret
        for vvv in d:
            ret[vvv]+=d[vvv]
    ret[c[v]]=size
    return size,ret

n=int(input())
c=list(map(int,input().split()))
g=[[] for _ in range(n+1)]
ans=[0]*(n+1)
for _ in range(n-1):
    s,t=map(int, input().split())
    s-=1
    t-=1
    g[s].append(t)
    g[t].append(s)
_,ret=dfs(0,-1)
for i in range(1,n+1):
    print(sum(n)-ans[i]-sum(n-ret[i]))

Solution (Python):
from collections import defaultdict
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)

def merge(a,b):
    if len(a) > len(b):
        a,b = b,a
    for k,v in a.items():
        b[k] += v
    return b

dp = [defaultdict(int) for _ in range(N)] 
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

        dp[p] = merge(dp[c],dp[p])
        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')

Solution (Python):
from collections import defaultdict
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)

def merge(a,b):
    if len(a) > len(b):
        a,b = b,a
    for k,v in a.items():
        b[k] += v
    return b

dp = [defaultdict(int) for _ in range(N)] 
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

        dp[p] = merge(dp[c],dp[p])
        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')

Solution (Python):
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 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


def tramp(gen):
    stack = [gen]
    val = None
    while stack:
        try:
            stack.append(stack[-1].send(val))
            val = None
        except StopIteration as e:
            stack.pop()
            val = e.value
    return val


_, root_clust = tramp(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)

Solution (Python):
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


def tramp(gen_fn, *args):
    stack = [gen_fn(*args)]
    val = None
    while stack:
        gen = stack[-1]
        try:
            params = gen.send(val)
        except StopIteration as e:
            stack.pop()
            val = e.value
        else:
            stack.append(go(*params))
            val = None
    return val


_, root_clust = tramp(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)

Solution (Python):
import sys
from collections import defaultdict

sys.setrecursionlimit(10 ** 6)


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

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().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)  
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

print(*answer[1:], sep='\n')

Solution (Python):
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)

for v in range(n):
    if idx[v] < len(g[v]) and 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):
    
    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
            
            p.append(v)
        else:
            p.append(v)
    q = n - acc[-1]
    
    ans[col] -= q*(q+1)//2
print(*ans[1:], sep='\n')

Solution (Python):
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: {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()))

Solution (Python):
import sys
input = lambda: sys.stdin.readline().rstrip()
N = int(input())
C = [int(a) - 1 for a in input().split()]
X = [[] for i in range(N)]
for i in range(N-1):
    x, y = map(int, input().split())
    X[x-1].append(y-1)
    X[y-1].append(x-1)

def EulerTour(n, X, i0):
    f = lambda k: k * (k + 1) // 2
    USED = [0] * n
    ORG = [0] * n
    TMP = [0] * n
    P = [-1] * n
    Q = [~i0, i0]
    ct = -1
    ET1 = [0] * n
    ET2 = [0] * n
    ANS = [f(n)] * n
    while Q:
        i = Q.pop()
        if i < 0:
            i = ~i
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                p = P[i]
                k = ET2[i] - ET1[i] + 1 - USED[C[p]] + ORG[i]
                ANS[C[p]] -= f(k)
                TMP[p] += k
            continue
        if i >= 0:
            if i: ORG[i] = USED[C[P[i]]]
            ct += 1
            if ET1[i] == 0: ET1[i] = ct
        for a in X[i][::-1]:
            if a != P[i]:
                P[a] = i
                for k in range(len(X[a])):
                    if X[a][k] == i:
                        del X[a][k]
                        break
                Q.append(~a)
                Q.append(a)
    for i in range(n):
        ANS[i] -= f(n - USED[i])
    return ANS

print(*EulerTour(N, X, 0), sep = "\n")

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().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)  
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(v, p):
    A = dp[v]
    B = dp[p]
    size_A = size[v]
    size_B = size[p]
    c = C[p]
    if len(A) < len(B):
        A, B = B, A
        size_A, size_B = size_B, size_A
    answer[c] += (size_A - A[c]) * (size_B - B[c])
    
    for key, cnt in B.items():
        answer[key] -= cnt * A[key]
        A[key] += cnt
    size[p] = size_A + size_B
    dp[p] = A


answer = [0] * (N + 1)
dp = [defaultdict(int) for _ in range(N + 1)]
size = [0] * (N + 1)

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

A = dp[root]
c = size[root]
for v in range(1, N + 1):
    print(c * A[v] + answer[v])

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().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)  
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(v, p):
    A = dp[v]
    B = dp[p]
    size_A = size[v]
    size_B = size[p]
    c = C[p]
    if len(A) < len(B):
        A, B = B, A
        size_A, size_B = size_B, size_A
    answer[c] += (size_A - A[c]) * (size_B - B[c])
    
    for key, cnt in B.items():
        answer[key] -= cnt * A[key]
        A[key] += cnt
    size[p] = size_A + size_B
    dp[p] = A


answer = [0] * (N + 1)
dp = [defaultdict(int) for _ in range(N + 1)]
size = [0] * (N + 1)

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

A = dp[root]
c = size[root]
for v in range(1, N + 1):
    print(c * A[v] + answer[v])

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().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)  
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(v, p):
    A = dp[v]
    B = dp[p]
    size_A = size[v]
    size_B = size[p]
    c = C[p]
    if len(A) < len(B):
        A, B = B, A
        size_A, size_B = size_B, size_A
    answer[c] += (size_A - A[c]) * (size_B - B[c])
    
    for key, cnt in B.items():
        answer[key] -= cnt * A[key]
        A[key] += cnt
    size[p] = size_A + size_B
    dp[p] = A


answer = [0] * (N + 1)
dp = [defaultdict(int) for _ in range(N + 1)]
size = [0] * (N + 1)

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

A = dp[root]
c = size[root]
for v in range(1, N + 1):
    print(c * A[v] + answer[v])

Solution (Python):
import sys
input = lambda: sys.stdin.readline().rstrip()
N = int(input())
C = [int(a) - 1 for a in input().split()]
X = [[] for i in range(N)]
for i in range(N-1):
    x, y = map(int, input().split())
    X[x-1].append(y-1)
    X[y-1].append(x-1)

def EulerTour(n, X, i0):
    f = lambda k: k * (k + 1) // 2
    USED = [0] * n
    ORG = [0] * n
    TMP = [0] * n
    P = [-1] * n
    Q = [~i0, i0]
    ct = -1
    ET1 = [0] * n
    ET2 = [0] * n
    ANS = [f(n)] * n
    while Q:
        i = Q.pop()
        if i < 0:
            i = ~i
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                p = P[i]
                k = ET2[i] - ET1[i] + 1 - USED[C[p]] + ORG[i]
                ANS[C[p]] -= f(k)
                TMP[p] += k
            continue
        if i >= 0:
            if i: ORG[i] = USED[C[P[i]]]
            ct += 1
            if ET1[i] == 0: ET1[i] = ct
        for a in X[i][::-1]:
            if a != P[i]:
                P[a] = i
                for k in range(len(X[a])):
                    if X[a][k] == i:
                        del X[a][k]
                        break
                Q.append(~a)
                Q.append(a)
    for i in range(n):
        ANS[i] -= f(n - USED[i])
    return ANS

print(*EulerTour(N, X, 0), sep = "\n")

Solution (Python):
import sys
input = lambda: sys.stdin.readline().rstrip()
N = int(input())
C = [int(a) - 1 for a in input().split()]
X = [[] for i in range(N)]
for i in range(N-1):
    x, y = map(int, input().split())
    x, y = x-1, y-1
    X[x].append(y)
    X[y].append(x)

def EulerTour(n, X, i0):
    def f(k):
        return k * (k + 1) // 2
    USED = [0] * n
    USEDORG = [0] * n
    USEDTMP = [0] * n
    P = [-1] * n
    Q = [~i0, i0]
    ct = -1
    ET = []
    ET1 = [0] * n
    ET2 = [0] * n
    DE = [0] * n
    ANS = [0] * n
    de = -1
    while Q:
        i = Q.pop()
        if i < 0:
            ET2[~i] = ct
            de -= 1
            USED[C[~i]] += 1 + USEDTMP[~i]
            USEDTMP[~i] = 0
            if ~i:
                p = P[~i]
                pc = C[p]
                k = ET2[~i] - ET1[~i] + 1 - USED[pc] + USEDORG[~i]
                ANS[pc] += f(k)
                USEDTMP[p] += k
            continue
        if i >= 0:
            if i:
                USEDORG[i] = USED[C[P[i]]]
            ET.append(i)
            ct += 1
            if ET1[i] == 0: ET1[i] = ct
            de += 1
            DE[i] = de
        for a in X[i][::-1]:
            if a != P[i]:
                P[a] = i
                for k in range(len(X[a])):
                    if X[a][k] == i:
                        del X[a][k]
                        break
                Q.append(~a)
                Q.append(a)
    for i in range(n):
        ANS[i] = f(n) - f(n - USED[i]) - ANS[i]
    return ANS

ANS = EulerTour(N, X, 0)
print(*ANS, sep = "\n")

Solution (Python):
import sys

sys.setrecursionlimit(2 * 10 ** 5)

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


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


rootTree(1)




colorlessCount = [0 for c in range(N + 1)]


colorlessTime = [0 for c in range(N + 1)]
t = 0


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


buildComps(1)


for c in range(1, N + 1):
    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])

Solution (Python):
import sys
from collections import defaultdict

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 = [0 for c in range(N + 1)]
colorlessTime = [0 for c in range(N + 1)]
t = 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)

for c in range(1, N + 1):
    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])

Solution (Python):
import os
from collections import Counter, defaultdict

import sys

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7


N = int(sys.stdin.buffer.readline())
C = list(map(int, sys.stdin.buffer.readline().split()))
AB = [list(map(int, sys.stdin.buffer.readline().split())) for _ in range(N - 1)]

graph = [[] for _ in range(N)]
for a, b in AB:
    a -= 1
    b -= 1
    graph[a].append(b)
    graph[b].append(a)

color_counts = [defaultdict(int) for _ in range(N + 1)]
for c in range(1, N + 1):
    color_counts[c][0] = N
color_v_weights = [[] for _ in range(N + 1)]
color_v = [0 for _ in range(N + 1)]
color_v_next = [1 for _ in range(N + 1)]







def dfs(v, p):
    color = C[v]
    color_v_prev = color_v[color]

    children_cv = []
    ret = 1
    for u in graph[v]:
        if u == p:
            continue
        cv = color_v_next[color]
        children_cv.append(cv)
        color_v[color] = cv
        color_v_next[color] += 1
        cnt = dfs(u, v)
        ret += cnt
        color_counts[color][cv] += cnt
    color_v[color] = color_v_prev
    color_counts[color][color_v_prev] -= ret
    return ret


dfs(0, None)
counter = Counter(C)

for color, cc in enumerate(color_counts[1:], start=1):
    ans = N * (N - 1) // 2 + counter[color]
    for cnt in cc.values():
        ans -= cnt * (cnt - 1) // 2
    print(ans)

Solution (Python):
import sys
sys.setrecursionlimit(10**9)

mod = 1000000007
INF = 110110110


N = int(input())
c = list(map(int, input().split()))
for i in range(N):
    c[i] -= 1

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



child = [0] * N

def dfs(now, par):
    res = 1
    for to in E[now]:
        if to == par:
            continue
        res += dfs(to, now)
    child[now] = res
    return res

dfs(0, -1)

ans = [0] * N
ok = [N] * N
def solve(now, par):
    memo = ok[c[now]]
    up = ok[c[now]] - child[now]
    num = [up]
    for to in E[now]:
        if to == par:
            continue
        num.append(child[to])
    S = sum(num)
    ans[c[now]] += S + 1
    for n in num:
        S -= n
        ans[c[now]] += S * n

    for to in E[now]:
        if to == par:
            continue
        ok[c[now]] = child[to]
        solve(to, now)
    ok[c[now]] = memo - child[now]

solve(0, -1)

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