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
sys.setrecursionlimit(10**7)

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

subtree=[1]*N
def size(v,pv):
    for nv in edge[v]:
        if nv!=pv:
            size(nv,v)
            subtree[v]+=subtree[nv]

size(0,-1)

data=[[] for i in range(N)]
Parent=[[0] for i in range(N)]
parent=[0]*N
def dfs(v,pv,nop):
    pp=parent[nop]
    parent[nop]=v
    Parent[nop].append(v)
    data[c[v]].append((parent[c[v]],v))
    for nv in edge[v]:
        if nv!=pv:
            dfs(nv,v,c[v])
    parent[nop]=pp

dfs(0,-1,0)

for i in range(N):
    dic={v:subtree[v] for v in Parent[i]}
    for p,ch in data[i]:
        dic[p]-=subtree[ch]

    res=N*(N+1)//2
    for p in dic:
        res-=dic[p]*(dic[p]+1)//2
    print(res)

Solution (Python):
import sys

input=sys.stdin.readline
sys.setrecursionlimit(10**7)

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

subtree=[1]*N
def size(v,pv):
    for nv in edge[v]:
        if nv!=pv:
            size(nv,v)
            subtree[v]+=subtree[nv]

size(0,-1)

data=[[] for i in range(N)]
Parent=[[0] for i in range(N)]
parent=[0]*N
def dfs(v,pv,nop):
    pp=parent[nop]
    parent[nop]=v
    Parent[nop].append(v)
    data[c[v]].append((parent[c[v]],v))
    for nv in edge[v]:
        if nv!=pv:
            dfs(nv,v,c[v])
    parent[nop]=pp

dfs(0,-1,0)

for i in range(N):
    dic={v:subtree[v] for v in Parent[i]}
    for p,ch in data[i]:
        dic[p]-=subtree[ch]

    res=N*(N+1)//2
    for p in dic:
        res-=dic[p]*(dic[p]+1)//2
    print(res)

Solution (Python):
import sys






def solve():
    n = int(input())
    c = list(map(lambda x : int(x)-1, input().split()))
    
    
    adj = [[] for i in range(n)]
    for i in range(n-1):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        adj[a].append(b)
        adj[b].append(a)

    
    stack = [[] for color in range(n)]
    
    
    size = [1]*n
    
    in_time = [-1]*n
    out_time = [-1]*n
    timer = 0
    total = n*(n+1)//2
    
    ans = [total]*n
    

    def dfs(parent, root):
        nonlocal timer
        in_time[root] = timer
        timer += 1

        for child in adj[root]:
            if parent == child:
                continue
            dfs(root, child)
            size[root] += size[child]

            cnt = size[child]
            while stack[c[root]]:
                x = stack[c[root]][-1]
                if in_time[x] > in_time[root] and out_time[x] != -1:
                    cnt -= size[x]
                    stack[c[root]].pop()
                else:
                    break
            ans[c[root]] -= cnt*(cnt+1)//2

        out_time[root] = timer
        timer += 1

        stack[c[root]].append(root)

    sys.setrecursionlimit(10**6)

    dfs(0, 0)
    
    for color in range(n):
        cnt = n
        while len(stack[color]) > 0:
            x = stack[color][-1]
            cnt -= size[x]
            stack[color].pop()
        
        ans[color] -= cnt*(cnt+1)//2
        print(ans[color])


solve()

Solution (Python):
import sys


def input():
    return sys.stdin.readline().rstrip()


def solve():
    n = int(input())
    c = list(map(lambda x : int(x)-1, input().split()))
    
    
    adj = [[] for i in range(n)]
    for i in range(n-1):
        a, b = map(int, input().split())
        a -= 1
        b -= 1
        adj[a].append(b)
        adj[b].append(a)

    
    stack = [[] for color in range(n)]
    
    
    size = [1]*n
    
    in_time = [-1]*n
    out_time = [-1]*n
    timer = 0
    total = n*(n+1)//2
    
    ans = [total]*n
    

    def dfs(parent, root):
        nonlocal timer
        in_time[root] = timer
        timer += 1

        for child in adj[root]:
            if parent == child:
                continue
            dfs(root, child)
            size[root] += size[child]

            cnt = size[child]
            while stack[c[root]]:
                x = stack[c[root]][-1]
                if in_time[x] > in_time[root] and out_time[x] != -1:
                    cnt -= size[x]
                    stack[c[root]].pop()
                else:
                    break
            ans[c[root]] -= cnt*(cnt+1)//2

        out_time[root] = timer
        timer += 1

        stack[c[root]].append(root)

    sys.setrecursionlimit(10**6)

    dfs(0, 0)
    
    for color in range(n):
        cnt = n
        while len(stack[color]) > 0:
            x = stack[color][-1]
            cnt -= size[x]
            stack[color].pop()
        
        ans[color] -= cnt*(cnt+1)//2
        print(ans[color])


solve()

Solution (Python):
import sys


def input(): return sys.stdin.readline().rstrip()


n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size = [1 for i in range(n)]
stack = [[] for color in range(n)]
in_time = [-1 for i in range(n)]
out_time = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]



def dfs(parent, root):
    global timer
    in_time[root] = timer
    timer += 1

    for child in adj[root]:
        if parent == child:
            continue
        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while stack[c[root]]:
            x = stack[c[root]][-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                stack[c[root]].pop()
            else:
                break
        ans[c[root]] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
from typing import List
import sys

input = lambda: sys.stdin.readline().rstrip()
n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size: List[int] = [1 for i in range(n)]
stack: List[List[int]] = [[] for color in range(n)]
in_time: List[int] = [-1 for i in range(n)]
out_time: List[int] = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans: List[int] = [total for color in range(n)]



def dfs(parent: int, root: int):
    global timer
    in_time[root] = timer
    timer += 1

    for child in adj[root]:
        if parent == child:
            continue
        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while stack[c[root]]:
            x = stack[c[root]][-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                stack[c[root]].pop()
            else:
                break
        ans[c[root]] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
from typing import List
import sys

n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size: List[int] = [1 for i in range(n)]
stack: List[List[int]] = [[] for color in range(n)]
in_time: List[int] = [-1 for i in range(n)]
out_time: List[int] = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans: List[int] = [total for color in range(n)]



def dfs(parent: int, root: int):
    global timer
    in_time[root] = timer
    timer += 1

    for child in adj[root]:
        if parent == child:
            continue
        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while stack[c[root]]:
            x = stack[c[root]][-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                stack[c[root]].pop()
            else:
                break
        ans[c[root]] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
import sys
n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size = [1 for i in range(n)]
stack = [[] for color in range(n)]
in_time = [-1 for i in range(n)]
out_time = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]



def dfs(parent, root):
    global timer
    in_time[root] = timer
    timer += 1

    for child in adj[root]:
        if parent == child:
            continue
        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while stack[c[root]]:
            x = stack[c[root]][-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                stack[c[root]].pop()
            else:
                break
        ans[c[root]] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
import sys
n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size = [1 for i in range(n)]
stack = [[] for color in range(n)]
in_time = [-1 for i in range(n)]
out_time = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]



def dfs(parent, root):
    global timer
    in_time[root] = timer
    timer += 1

    color = c[root]
    s = stack[color]
    for child in adj[root]:
        if parent == child:
            continue

        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while s:
            x = s[-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                s.pop()
            else:
                break
        ans[color] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
import sys
n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size = [1 for i in range(n)]
stack = [[] for color in range(n)]
in_time = [-1 for i in range(n)]
out_time = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]



def dfs(parent, root):
    global timer
    in_time[root] = timer
    timer += 1

    for child in adj[root]:
        if parent == child:
            continue
        dfs(root, child)
        size[root] += size[child]

        color = c[root]
        cnt = size[child]
        s = stack[c[root]]
        while s:
            x = s[-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                s.pop()
            else:
                break
        ans[color] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
import sys
n = int(input())
c = list(map(int, input().split()))
c = [x-1 for x in c]

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


size = [1 for i in range(n)]
stack = [[] for color in range(n)]
in_time = [-1 for i in range(n)]
out_time = [-1 for i in range(n)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]



def dfs(parent, root):
    global timer
    in_time[root] = timer
    timer += 1

    for child in adj[root]:
        if parent == child:
            continue
        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while stack[c[root]]:
            x = stack[c[root]][-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                stack[c[root]].pop()
            else:
                break
        ans[c[root]] -= cnt*(cnt+1)//2

    out_time[root] = timer
    timer += 1

    stack[c[root]].append(root)


sys.setrecursionlimit(10**6)

dfs(0, 0)

for color in range(n):
    cnt = n
    while len(stack[color]) > 0:
        x = stack[color][-1]
        cnt -= size[x]
        stack[color].pop()
    
    ans[color] -= cnt*(cnt+1)//2
    print(ans[color])

Solution (Python):
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**9)

def solve():
    N = int(input())
    Cs = [A-1 for A in map(int, input().split())]
    adjL = [[] for _ in range(N)]
    for _ in range(N-1):
        a, b = map(int, input().split())
        a, b = a-1, b-1
        adjL[a].append(b)
        adjL[b].append(a)

    sizeSubtrees = [1] * N
    numVs = [0] * N
    anss = [0] * N

    def dfs(vNow, vPar, clrPar):
        clrNow = Cs[vNow]
        numVNow = numVs[clrNow]
        numVPar = numVs[clrPar]
        for v2 in adjL[vNow]:
            if v2 == vPar: continue
            dfs(v2, vNow, clrNow)
            sizeSubtrees[vNow] += sizeSubtrees[v2]
        numVs[clrNow] = numVNow + sizeSubtrees[vNow]
        if clrPar != -1:
            k = sizeSubtrees[vNow] - (numVs[clrPar]-numVPar)
            anss[clrPar] += k*(k+1)//2

    dfs(0, -1, -1)

    for clr in range(N):
        if clr == Cs[0]: continue
        k = N - numVs[clr]
        anss[clr] += k*(k+1)//2

    numAll = N*(N+1)//2
    anss = [numAll-ans for ans in anss]

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


solve()

Solution (Python):
def f(l):
    r = sum(l)**2
    for x in l:
        r -= x**2
    r = r//2
    r += sum(l)+1
    return r

N, = map(int, input().split())
G = [set() for _ in range(N+1)]
G2 = [set() for _ in range(N+1)]
sz = [1 for _ in range(N+1)]
X = [-1] + list(map(int, input().split()))
for _ in range(N-1):
    a, b = map(int, input().split())
    G[a].add(b)
    G[b].add(a)

stack = [1]
vs = set([1])
vs_dfs = []
while stack:
    v = stack.pop()
    vs_dfs.append(v)
    for u in G[v]:
        if u in vs:
            continue
        vs.add(u)
        G2[v].add(u)
        stack += [u]
for x in vs_dfs[::-1]:
    for u in G2[x]:
        sz[x] += sz[u]

stack = [1]
path = []
R = [0] * (N+1)
kh = [[N] for _ in range(N+1)]
while stack:
    v = stack.pop()
    if v > 0:
        if path:
            parent = path[-1]
            kh[X[parent]].append(sz[v])
        path.append(v)
        tmp = []
        for u in G2[v]:
            tmp.append(sz[u])
            stack += [-v, u]
        tmp.append(kh[X[v]][-1]-1-sum(tmp))
        kh[X[v]][-1] -= sz[v]
        R[X[v]] += f(tmp)
    else:
        c = path.pop()
        kh[X[-v]].pop()
for i in range(1, N+1):
    print(R[i])

Solution (Python):
def f(l):
    r = sum(l)**2
    for x in l:
        r -= x**2
    r = r//2
    r += sum(l)+1
    return r

N, = map(int, input().split())
G = [set() for _ in range(N+1)]
G2 = [set() for _ in range(N+1)]
sz = [1 for _ in range(N+1)]
X = [-1] + list(map(int, input().split()))
for _ in range(N-1):
    a, b = map(int, input().split())
    G[a].add(b)
    G[b].add(a)

stack = [1]
vs = set([1])
vs_dfs = []
while stack:
    v = stack.pop()
    vs_dfs.append(v)
    for u in G[v]:
        if u in vs:
            continue
        vs.add(u)
        G2[v].add(u)
        stack += [u]
G = None

for x in vs_dfs[::-1]:
    for u in G2[x]:
        sz[x] += sz[u]

stack = [1]
path = []
R = [0] * (N+1)
kh = [[] for _ in range(N+1)]
tt = [N-1 for _ in range(N+1)]
while stack:
    v = stack.pop()
    if v > 0:
        if path:
            parent = path[-1]
            kh[X[parent]].append(sz[v])
        path.append(v)
        tmp = []
        for u in G2[v]:
            tmp.append(sz[u])
            stack += [-v, u]
        if kh[X[v]]:
            tmp.append(kh[X[v]][-1]-1-sum(tmp))
            kh[X[v]][-1] -= sz[v]
        else:
            tmp.append(tt[X[v]]-sum(tmp))
            tt[X[v]] -= sz[v]
        R[X[v]] += f(tmp)

    else:

        c = path.pop()
        kh[X[-v]].pop()
for i in range(1, N+1):
    print(R[i])

Solution (Python):
import bisect
import collections
import sys

sys.setrecursionlimit(1000000)
input = sys.stdin.readline
ACMOD = 1000000007
INF = 1 << 62


def lmi():
    return list(map(int, input().split()))


def llmi(n):
    return [lmi() for _ in range(n)]


N = int(input())
c_raw = lmi()
ab = llmi(N - 1)

adj = collections.defaultdict(list)
for a, b in ab:
    adj[a - 1].append(b - 1)
    adj[b - 1].append(a - 1)

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

visited = [False] * N

visited[0] = True
stacks = [[[0, 0]] for i in range(N)]
ans = [0] * N
curr = 1














color_stack = [[0, N] for _ in range(N)]



class Tree:
    def __init__(self, i, _c):
        self.i = i
        self.color = _c
        self.child_count = 0
        self.children = []

    def set_child_count(self, parent=-1):
        for node_index in adj[self.i]:
            if node_index == parent:
                continue
            self.children.append(node_index)
            trees[node_index].set_child_count(self.i)
            self.child_count += trees[node_index].child_count + 1

    def dfs(self):
        color_stack[self.color][-1] -= self.child_count + 1
        for node_index in self.children:
            color_stack[self.color].append(trees[node_index].child_count + 1)
            trees[node_index].dfs()
            ans[self.color] -= get_path_count(color_stack[self.color].pop())
            


trees = [Tree(i, c - 1) for i, c in enumerate(c_raw)]


def get_path_count(x):
    return x + x * (x - 1) // 2


ans = [get_path_count(N)] * N
trees[0].set_child_count()
trees[0].dfs()

for color in range(N):
    
    
    print(ans[color]-get_path_count(color_stack[color][-1]))

Solution (Python):
import bisect
import collections
import sys

sys.setrecursionlimit(1000000)
input = sys.stdin.readline
ACMOD = 1000000007
INF = 1 << 62


def lmi():
    return list(map(int, input().split()))


def llmi(n):
    return [lmi() for _ in range(n)]


N = int(input())
c_raw = lmi()
ab = llmi(N - 1)

adj = collections.defaultdict(list)
for a, b in ab:
    adj[a - 1].append(b - 1)
    adj[b - 1].append(a - 1)

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

visited = [False] * N

visited[0] = True
stacks = [[[0, 0]] for i in range(N)]
ans = [0] * N
curr = 1














color_stack = [[0, N] for _ in range(N)]



class Tree:
    def __init__(self, i, _c):
        self.i = i
        self.color = _c
        self.child_count = 0
        self.children = []

    def set_child_count(self, parent=-1):
        for node_index in adj[self.i]:
            if node_index == parent:
                continue
            self.children.append(node_index)
            trees[node_index].set_child_count(self.i)
            self.child_count += trees[node_index].child_count + 1

    def dfs(self):
        color_stack[self.color][-1] -= self.child_count + 1
        for node_index in self.children:
            color_stack[self.color].append(trees[node_index].child_count + 1)
            trees[node_index].dfs()
            ans[self.color] -= get_path_count(color_stack[self.color].pop())
            


trees = [Tree(i, c - 1) for i, c in enumerate(c_raw)]


def get_path_count(x):
    return x + x * (x - 1) // 2


ans = [get_path_count(N)] * N
trees[0].set_child_count()
trees[0].dfs()

for color in range(N):
    
    
    print(ans[color]-get_path_count(color_stack[color][-1]))

Solution (Python):
import bisect
import collections
import sys

sys.setrecursionlimit(1000000)
input = sys.stdin.readline
ACMOD = 1000000007
INF = 1 << 62


def lmi():
    return list(map(int, input().split()))


def llmi(n):
    return [lmi() for _ in range(n)]


N = int(input())
c_raw = lmi()
ab = llmi(N - 1)

adj = collections.defaultdict(list)
for a, b in ab:
    adj[a - 1].append(b - 1)
    adj[b - 1].append(a - 1)

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

visited = [False] * N

visited[0] = True
stacks = [[[0, 0]] for i in range(N)]
ans = [0] * N
curr = 1














color_stack = [[0, N] for _ in range(N)]
color_groups = [[] for _ in range(N)]


class Tree:
    def __init__(self, i, _c):
        self.i = i
        self.color = _c
        self.child_count = 0
        self.children = []

    def set_child_count(self, parent=-1):
        for node_index in adj[self.i]:
            if node_index == parent:
                continue
            self.children.append(node_index)
            trees[node_index].set_child_count(self.i)
            self.child_count += trees[node_index].child_count + 1

    def dfs(self):
        color_stack[self.color][-1] -= self.child_count + 1
        for node_index in self.children:
            color_stack[self.color].append(trees[node_index].child_count + 1)
            trees[node_index].dfs()
            color_groups[self.color].append(color_stack[self.color].pop())


trees = [Tree(i, c - 1) for i, c in enumerate(c_raw)]
trees[0].set_child_count()
trees[0].dfs()

for color in range(N):
    v = color_groups[color] + [color_stack[color][-1]]
    print(N + N * (N - 1) // 2 - sum(map(lambda x: x + x * (x - 1) // 2, v)))

Solution (Python):



import sys
sys.setrecursionlimit(10**7)
from pprint import pprint as pp
from pprint import pformat as pf



import math

import bisect

class BIT: 
    """
    http://hos.ac/slides/20140319_bit.pdf
    """

    NORMAL = 0
    ROLLBACK = 1

    def __init__(self, size):
        self.size = size
        self.data = [0] * (size + 1) 
        self.status = BIT.NORMAL
        self.history = []

    def add(self, pos, val):
        assert pos > 0, pos
        if self.status == BIT.NORMAL:
            self.history.append((pos, val))
        k = pos
        while k <= self.size:
            self.data[k] += val
            
            k += k & -k

    def sum(self, pos):
        s = 0
        k = pos
        while k > 0:
            s += self.data[k]
            
            k -= k & -k
        return s

    def sum_section(self, frm, to):
        to_val = self.sum(to)
        frm_val = self.sum(frm)
        return to_val - frm_val

    def save(self):
        self.history = []

    def rollback(self):
        self.status = BIT.ROLLBACK
        for (pos, val) in self.history:
            self.add(pos, -1 * val)
        self.history = []
        self.status = BIT.NORMAL

class Graph:

    def __init__(self, size):
        self.size = size
        self.vertices = [0] * size 
        self.edges = [None] * size 
        for i in range(size):
            self.edges[i] = []

    def add_edge(self, frm, to):
        self.edges[frm].append(to)
        self.edges[to].append(frm)

def make_mybit(size):
    bit = BIT(size)
    for i in range(size):
        bit.add(i + 1, 1)
    bit.save()
    return bit

def init_color_s(size):
    color_s = [None] * graph.size
    for i in range(graph.size):
        color_s[i] = []
    return color_s

def dfs(graph):
    parent_s = [0] * graph.size
    in_s = [0] * graph.size
    out_s = [0] * graph.size
    order = 0 
    color_s = init_color_s(graph.size)
    def _dfs(v, prev):
        nonlocal order
        order += 1
        parent_s[v] = prev
        in_s[v] = order
        c = graph.vertices[v]
        color_s[c].append(v)
        for to in graph.edges[v]:
            if to == prev:
                continue
            _dfs(to, v)
        out_s[v] = order
    _dfs(0, -1)
    return parent_s, in_s, out_s, color_s

def calc(v):
    return v * (v + 1) // 2

def get_connected(bit, v, in_s, out_s):
    
    
    large_key = out_s[v]
    small_key = in_s[v] - 1
    return bit.sum_section(small_key, large_key)

def calc_for(bit, v, in_s, out_s):
    connected = get_connected(bit, v, in_s, out_s)
    
    
    return calc(connected)

def cut(bit, v, in_s, out_s):
    connected = get_connected(bit, v, in_s, out_s)
    bit.add(in_s[v], -1 * connected)

def solve(n, graph):
    bit = make_mybit(graph.size)
    parent_s, in_s, out_s, color_s = dfs(graph)
    
    
    
    
    for same_color_s in color_s:
        ans = calc(graph.size)
        bit.rollback()
        
        
        for v in same_color_s[::-1]:
            for to in graph.edges[v]:
                if to == parent_s[v]:
                    continue
                ans -= calc_for(bit, to, in_s, out_s)
            cut(bit, v, in_s, out_s)
        ans -= calc_for(bit, 0, in_s, out_s)
        
        print(ans)






if __name__ == '__main__':
    n = int(input())
    graph = Graph(n)
    vertices = list(map(int, input().split()))
    for i, c in enumerate(vertices):
        vertices[i] = c - 1
    graph.vertices = vertices
    for _ in range(n - 1):
        frm, to = list(map(int, input().split()))
        frm -= 1
        to -= 1
        graph.add_edge(frm, to)
    solve(n, graph)

Solution (Python):
import sys
sys.setrecursionlimit(10**6)
 
N = int(input())  
C = [int(x) - 1  for x in input().split()]  
 
E = [[] for _ in range(N)]
for _ in range(N-1):
    a, b = map(int, input().split())
    E[a-1].append(b-1)
    E[b-1].append(a-1)
    
def dfs(u, p = -1):
 
    color = C[u]  
    VC[color].append(u)  
    
    s = 1  
    for v in E[u]:
        if v == p:
            continue
            
        child[u] = 0  
        ret = dfs(v, u)
 
        s += ret  
        child[u] += ret  
        
        ans[color] -= child[u] * (child[u] + 1) // 2        
 
    VC[color].pop()
    if VC[color]: 
        child[VC[color][-1]] -= s  
    else:  
        root_size[color] -= s  
        
    return s

child = [0] * N  
root_size = [N] * N  
VC = [ [] for _ in range(N)]  
ans = [N*(N+1)//2] * N  
 
dfs(0)
 
for i in range(N):
    print(ans[i] - root_size[i] * (root_size[i] + 1) // 2)

Solution (Python):
import sys
sys.setrecursionlimit(10**8)
input = sys.stdin.readline
N = int(input())
C = list(map(int,input().split()))
AB = [tuple(map(int,input().split())) for i in range(N-1)]

es = [[] for _ in range(N)]
for a,b in AB:
    a,b = a-1,b-1
    es[a].append(b)
    es[b].append(a)
cs = [[] for _ in range(N)]
for i,c in enumerate(C):
    cs[c-1].append(i)

tin = [-1]*N
tout = [-1]*N
k = 0
def dfs(v,p=-1):
    global k
    tin[v] = k
    k += 1
    for to in es[v]:
        if to==p: continue
        dfs(to,v)
    tout[v] = k
dfs(0)

class BinaryIndexedTree:
    def __init__(self,size):
        self.N = size
        self.bit = [0]*(size+1)
    def add(self,x,w): 
        x += 1
        while x <= self.N:
            self.bit[x] += w
            x += (x & -x)
    def _sum(self,x): 
        ret = 0
        while x > 0:
            ret += self.bit[x]
            x -= (x & -x)
        return ret
    def sum(self,l,r): 
        return self._sum(r) - self._sum(l)
    def __str__(self): 
        arr = [self.sum(i,i+1) for i in range(self.N)]
        return str(arr)
bit = BinaryIndexedTree(N)
for i in range(N):
    bit.add(i,1)

whole = N*(N+1)//2
anss = []
for i in range(N):
    ans = whole
    cs[i].sort(key=lambda x: -tin[x])
    history = []
    for v in cs[i]:
        cnt = 1
        for to in es[v]:
            if tin[to] < tin[v]: continue
            num = bit.sum(tin[to], tout[to])
            ans -= num*(num+1)//2
            cnt += num
        bit.add(tin[v], -cnt)
        history.append((tin[v], cnt))
    pn = bit.sum(0,N)
    ans -= pn*(pn+1)//2
    for a,b in history:
        bit.add(a,b)
    anss.append(ans)
print(*anss, sep='\n')