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
from collections import deque

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

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



deq = deque([1])
parent = [-1 for i in range(N+1)]
parent[1] = 0
topo = []
while deq:
    v = deq.popleft()
    topo.append(v)
    for nv in edge[v]:
        if parent[nv]==-1:
            parent[nv] = v
            deq.append(nv)

topo = topo[::-1]



vertex = []
size = [1 for i in range(N+1)]
size[0] = N

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

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

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

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

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

    return res[::-1]

ans = [0]*(N+1)

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

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

Solution (Python):


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

import numpy as np

def sort_edges(N, edges_):
    
    M = len(edges_)
    edges = np.empty((M * 2, 2), dtype=np.int64)
    edges[:M] = edges_
    edges[M:] = edges_[:, ::-1]
    order = np.argsort(edges[:, 0])  
    edges = edges[order, 1]
    c = np.zeros(N+1, dtype=np.int64)
    c_ = np.bincount(edges_.ravel())  
    c[:len(c_)] = c_
    c = np.cumsum(c)
    lefts = np.zeros(len(c) + 1, dtype=np.int64)
    lefts[1:] = c
    return edges, lefts

def eular_tour(edges, lefts, root):
    
    n = len(lefts)-1
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  
            continue
        p = parents[v]
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in edges[lefts[v]:lefts[v+1]]:
            if p != u:
                parents[u] = v
                stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

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

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

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

    return Ans[1:]


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


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

main()

Solution (Python):



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

import numpy as np

def eular_tour(G, root):
    
    
    n = len(G)
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  
            continue
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            parents[u] = v
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

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

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

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

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

    return Ans[1:]


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


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

main()

Solution (Python):



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

import numpy as np

def eular_tour(G, root):
    
    
    n = len(G)
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  
            continue
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            parents[u] = v
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

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

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

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

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

    return Ans[1:]


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


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

Solution (Python):



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

import numpy as np

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

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

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

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

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

Solution (Python):


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

import numpy as np

def eular_tour(G, root):
    
    
    n = len(G)
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  
            continue
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            parents[u] = v
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

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

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

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

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

    return Ans[1:]


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


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

Solution (Python):


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

import numpy as np

def eular_tour(G, root):
    
    
    n = len(G)
    stack = [root]
    tour = [0] * 0
    first = np.full(n, -100, dtype=np.int64)
    last = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if first[v] >= 0:
            last[v] = len(tour)
            tour.append(-v)  
            continue
        first[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, first, last

def dfs1(v, G, Siz):
    tour, first, last = eular_tour(G, 0)
    Siz[:] = last - first + 1 >> 1
    return 0









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

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

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

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

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

import numpy as np

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

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

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

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

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

Solution (Python):
import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = [int(c)-1 for c in input().split()]
G = [[] for _ in range(N)]
for _ in range(N-1):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    G[a].append(b)
    G[b].append(a)

Cnt = [0] * N
Ans = [0] * N
Siz = [0] * N

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



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

dfs1(0)
dfs(0)
print("\n".join(map(str, Ans)))

Solution (Python):
import sys
INF = 1 << 60
MOD = 10**9 + 7 
sys.setrecursionlimit(2147483647)
input = lambda:sys.stdin.readline().rstrip()

class SegmentTree(object):
    def __init__(self, A, dot, unit):
        n = 1 << (len(A) - 1).bit_length()
        tree = [unit] * (2 * n)
        for i, v in enumerate(A):
            tree[i + n] = v
        for i in range(n - 1, 0, -1):
            tree[i] = dot(tree[i << 1], tree[i << 1 | 1])
        self._n = n
        self._tree = tree
        self._dot = dot
        self._unit = unit

    def __getitem__(self, i):
        return self._tree[i + self._n]

    def update(self, i, v):
        i += self._n
        self._tree[i] = v
        while i != 1:
            i >>= 1
            self._tree[i] = self._dot(self._tree[i << 1], self._tree[i << 1 | 1])

    def add(self, i, v):
        self.update(i, self[i] + v)

    def sum(self, l, r):
        l += self._n
        r += self._n
        l_val = r_val = self._unit
        while l < r:
            if l & 1:
                l_val = self._dot(l_val, self._tree[l])
                l += 1
            if r & 1:
                r -= 1
                r_val = self._dot(self._tree[r], r_val)
            l >>= 1
            r >>= 1
        return self._dot(l_val, r_val)

from operator import add
def resolve():
    n = int(input())
    color = [[] for _ in range(n)]
    for v, c in enumerate(map(int, input().split())):
        color[c - 1].append(v)
    E = [[] for _ in range(n)]
    for _ in range(n - 1):
        u, v = map(int, input().split())
        u -= 1; v -= 1
        E[u].append(v)
        E[v].append(u)

    
    start = [-1] * n
    end = [-1] * n
    count = 0
    stack = [(~0, -1), (0, -1)]
    while stack:
        v, p = stack.pop()
        if v >= 0:
            start[v] = count
            count += 1
            for nv in E[v]:
                if nv == p:
                    continue
                stack.append((~nv, v))
                stack.append((nv, v))
        else:
            end[~v] = count

    f = lambda x : x * (x + 1) // 2
    ans = [f(n)] * n
    tree = SegmentTree([1] * n, add, 0)
    history = []
    for c in range(n):
        for v in sorted(color[c], key = lambda x : start[x], reverse = 1):
            count = 1
            for nv in E[v]:
                if start[v] > start[nv]: 
                    continue
                s = tree.sum(start[nv], end[nv])
                ans[c] -= f(s)
                count += s
            tree.add(start[v], -count)
            history.append((start[v], count))
        ans[c] -= f(tree.sum(0, n))
        while history:
            tree.add(*history.pop())

    print(*ans, sep = '\n')
resolve()

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

n = int(input())
c = list(map(lambda x: int(x) - 1, input().split()))

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

ans = [n * (n + 1) // 2] * n
n_subtree = [1] * n  
cnt = [0] * n  
n_visited = 0


def dfs(v, v_p):
    global n_visited
    cnt_v_before = cnt[c[v]]
    for v_next in graph[v]:
        if v_next == v_p:
            continue

        m_prev = n_visited - cnt[c[v]]
        dfs(v_next, v)
        n_subtree[v] += n_subtree[v_next]
        m_next = n_visited - cnt[c[v]]

        m = m_next - m_prev
        ans[c[v]] -= m * (m + 1) // 2

    cnt[c[v]] = cnt_v_before + n_subtree[v]
    n_visited += 1


dfs(0, -1)

for i in range(n):
    m = n - cnt[i]
    ans[i] -= m * (m + 1) // 2

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

Solution (Python):
def euler_tour(tree: list, root: int):
    """オイラーツアーを行う"""
    n = len(tree)
    par = {root: -1}
    trail = []
    is_up = []
    stack = [root]
    while stack:
        v = stack.pop()
        if v >= 0:  
            trail.append(v)
            is_up.append(0)
            stack.append(~v)
            for nxt_v in tree[v]:
                if nxt_v not in par:
                    par[nxt_v] = v
                    stack.append(nxt_v)
                    

        else:  
            if ~v == root:
                continue
            trail.append(par[~v])
            is_up.append(1)
 
    return trail, is_up


class BIT:
    """区間加算、一点取得クエリをそれぞれO(logN)で答えるデータ構造"""
    def __init__(self, n):
        self.size = n
        self.bit = [0] * (n + 1)

    def build(self, array):
        """arrayを初期値とするBinaryIndexTreeを構築する O(N)"""
        for i in range(1, self.size + 1):
            if i - 1 + (i & -i) >= self.size:
                self.bit[i] = array[i - 1]
            else:
                self.bit[i] = array[i - 1] - array[i + (i & -i) - 1]

    def _add(self, i, val):
        while i > 0:
            self.bit[i] += val
            i -= i & -i

    def get_val(self, i):
        """i番目の値を求める"""
        i = i + 1
        s = 0
        while i <= self.size:
            s += self.bit[i]
            i += i & -i
        return s

    def add(self, l, r, val):
        """区間[l, r)にvalを加える"""
        self._add(r, val)
        self._add(l, -val)


from collections import deque
import sys
input = sys.stdin.readline


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

tree = [[] for i in range(n)]
for a, b in info:
    a -= 1
    b -= 1
    tree[a].append(b)
    tree[b].append(a)


plus = n + 1
for i in range(1, n):
   if len(tree[i]) == 1:
        info.append([i + 1, plus])
        plus += 1

info.append([1, plus])

new_n = len(info) + 1
tree = [[] for i in range(new_n)]
for a, b in info:
    a -= 1
    b -= 1
    tree[a].append(b)
    tree[b].append(a)

root = new_n - 1
et, is_up = euler_tour(tree, root)

bit = BIT(n)
stack = [[] for i in range(n)]
ans = [0] * n

for i in range(len(et)):
    if et[i] >= n:
        continue

    if is_up[i]:
        ind = i
        i = c[et[i]] - 1
        val = bit.get_val(i)
        ans[i] += val * (val + 1) // 2

        bit.add(i, i + 1, -val)
        if len(stack[i]) >= 1 and is_up[ind + 1]:                                                               
            tmp = stack[i].pop()                                                             
            bit.add(i, i + 1, tmp)    
    else:
        i = c[et[i]] - 1
        val = bit.get_val(i)
        stack[i].append(val)
        bit.add(0, n, 1)
        bit.add(i, i + 1, -(val + 1))

for i in range(n):
    val = bit.get_val(i)
    ans[i] += val * (val + 1) // 2

for i in range(n):
    print(n * (n + 1) // 2 - ans[i])

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

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

    
    
    size_subtree = [1] * N
    num_v = [0] * N
    ans = [0] * N

    def dfs(v_now, v_parent, color_parent):
        color_now = C[v_now]
        num_v_now = num_v[color_now]
        num_v_parent = num_v[color_parent]
        for v2 in tree[v_now]:
            if v2 == v_parent: continue
            dfs(v2, v_now, color_now)
            size_subtree[v_now] += size_subtree[v2]
        num_v[color_now] = num_v_now + size_subtree[v_now]
        if color_parent != -1:
            k = size_subtree[v_now] - (num_v[color_parent]-num_v_parent)
            ans[color_parent] += k*(k+1)//2

    dfs(0, -1, -1)

    for color in range(N):
        if color == C[0]: continue
        k = N - num_v[color]
        ans[color] += k*(k+1)//2

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

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

main()