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 = 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
    ct = -1
    ET1 = [0] * n
    ET2 = [0] * n
    ANS = [f(n)] * n
    IND = [len(x) for x in X]
    i = i0
    while i >= 0:
        ind = IND[i]
        if not ind:
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                k = ET2[i] - ET1[i] + 1 - USED[C[P[i]]] + ORG[i]
                ANS[C[P[i]]] -= f(k)
                TMP[P[i]] += k
            i = P[i]
            continue
        
        j = X[i][ind - 1]
        if P[i] == j:
            IND[i] -= 1
            continue
        
        P[j] = i
        IND[i] -= 1
        i = j
        ORG[i] = USED[C[P[i]]]
        ct += 1
        ET1[i] = ct
    
    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[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
    ct = -1
    ET1 = [0] * n
    ET2 = [0] * n
    ANS = [f(n)] * n
    IND = [0] * n
    i = i0
    while i >= 0:
        ind = IND[i]
        if ind == 0:
            if i: ORG[i] = USED[C[P[i]]]
            ct += 1
            ET1[i] = ct
        
        if ind >= len(X[i]):
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                k = ET2[i] - ET1[i] + 1 - USED[C[P[i]]] + ORG[i]
                ANS[C[P[i]]] -= f(k)
                TMP[P[i]] += k
            i = P[i]
            continue
        
        j = X[i][ind]
        if P[i] == j:
            IND[i] += 1
            continue
        P[j] = i
        IND[i] += 1
        i = j
    
    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[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
    ct = -1
    ET1 = [0] * n
    ET2 = [0] * n
    ANS = [f(n)] * n
    IND = [0] * n
    i = i0
    while i >= 0:
        if IND[i] == 0:
            if i: ORG[i] = USED[C[P[i]]]
            ct += 1
            ET1[i] = ct
        
        if IND[i] >= len(X[i]):
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                k = ET2[i] - ET1[i] + 1 - USED[C[P[i]]] + ORG[i]
                ANS[C[P[i]]] -= f(k)
                TMP[P[i]] += k
            i = P[i]
            continue
        
        if P[i] == X[i][IND[i]]:
            IND[i] += 1
            continue
        P[X[i][IND[i]]] = i
        IND[i], i = IND[i] + 1, X[i][IND[i]]
    
    for i in range(n):
        ANS[i] -= f(n - USED[i])
    return ANS

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

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


n = int(input())
C = list(map(int,input().split()))





path = [set() for _ in range(n)]

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

ans = [n*(n+1)//2]*n
cparent = [[] for _ in range(n)]
root_size = [n]*n

reached = [False]*n
size = [0]*n


def dfs(p):
    c = C[p] - 1
    cparent[c].append(p)
    s = 1
    for nxt in path[p]:
        if reached[nxt]:continue
        reached[nxt] = True
        size[p] = 0
        ret = dfs(nxt)
        s += ret
        size[p] += ret
        ans[c] -= size[p] * (size[p] + 1)//2
    cparent[c].pop()
    
    if cparent[c]:
        size[cparent[c][-1]] -= s
    else:
        root_size[c] -= s
    
    return s

reached = [False]*n
reached[0] = True

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

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


def EulerTour(graph, root=1):
    global ind_L, ind_R, parent
    N = len(graph) - 1
    stack = [-root, root]
    tour = []
    i = -1
    while stack:
        v = stack.pop()
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
        yield v
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
parent = [0] * (N + 1)
ind_L = [0] * (N + 1)
ind_R = [0] * (N + 1)
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in EulerTour(graph):
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        c = C[parent[v]]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  
        x -= removed[c] - memo[v]
        answer[c] -= x * (x + 1) // 2
        removed[c] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))

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


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

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = [0] * (N + 1)
        ind_L = [0] * (N + 1)
        ind_R = [0] * (N + 1)
        i = -1
        tour = []
        stack = [-root, root]
        stack[0] = -root
        stack[1] = root
        while stack:
            v = stack.pop()
            tour.append(v)
            i += 1
            if v > 0:
                ind_L[v] = i
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack.append(-w)
                    stack.append(w)
                    k = nxt[k]
            else:
                ind_R[-v] = i
        return tour, ind_L, ind_R, parent
    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = [0] * (N + 1)
    answer = [N * (N + 1) // 2] * (N + 1)
    memo = [0] * (N + 1)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
AB = tuple(map(int, read().split()))
answer = solve(N, C, AB)
print('\n'.join(map(str, answer[1:])))

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


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

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = [0] * (N + 1)
        ind_L = [0] * (N + 1)
        ind_R = [0] * (N + 1)
        i = -1
        tour = []
        stack = [-root, root]
        stack[0] = -root
        stack[1] = root
        while stack:
            v = stack.pop()
            tour.append(v)
            i += 1
            if v > 0:
                ind_L[v] = i
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack.append(-w)
                    stack.append(w)
                    k = nxt[k]
            else:
                ind_R[-v] = i
        return tour, ind_L, ind_R, parent
    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = [0] * (N + 1)
    answer = [N * (N + 1) // 2] * (N + 1)
    memo = [0] * (N + 1)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
AB = tuple(map(int, read().split()))
answer = solve(N, C, AB)
print('\n'.join(map(str, answer[1:])))

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


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

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

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


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


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

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


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

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

    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)
    for v in seq:
        for u in child[v]:
            size[v] += size[u]

    ans = [calc(N)] * (N+1)
    par_C = [[] for _ in range(N+1)]
    sub = [0] * (N+1)
    sub_0 = [0] * (N+1)
    st = deque()
    st.append(1)
    while st:
        v = st.pop()
        if v > 0:
            st.append(-v)
            for u in child[v]:
                st.append(u)
            c = C[v-1]
            if par_C[c]:
                sub[par_C[c][-1]] += size[v]
            else:
                sub_0[c] += size[v]
            par_C[c].append(v)
        else:
            v = -v
            c = C[v-1]
            assert par_C[c][-1] == v
            par_C[c].pop()
            p = par[v]
            if p:
                cp = C[p-1]
                ans[cp] -= calc(size[v] - sub[p])
                sub[p] = 0
    for c in range(1, N+1):
        ans[c] -= calc(N - sub_0[c])
        print(ans[c])


if __name__ == '__main__':
    main()

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np
from numba import njit

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


@njit('(i4[:],i4[:],i4[:],i4[:],i4[:])', cache=True)
def main(C, tour, ind_L, ind_R, parent):
    N = len(C) - 1
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


@njit('(i4[:],i4[:],i4,i4)', cache=True)
def edges_to_array(A, B, N, M):
    head = np.full(N + 1, -1, np.int32)
    nxt = np.empty(M + M, np.int32)
    to = np.empty(M + M, np.int32)
    for i in range(M):
        a = A[i]
        b = B[i]
        nxt[i << 1] = head[a]
        to[i << 1] = b
        head[a] = i << 1
        nxt[i << 1 | 1] = head[b]
        to[i << 1 | 1] = a
        head[b] = i << 1 | 1
    return head, nxt, to


N = int(readline())
C = np.zeros(N + 1, np.int32)
C[1:] = np.array(readline().split(), np.int32)
AB = np.array(read().split(), np.int32)
head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)

tour, ind_L, ind_R, parent = EulerTour(head, nxt, to, 1)
answer = main(C, tour, ind_L, ind_R, parent)
print('\n'.join(answer[1:].astype(str)))

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
    done = [0] * n
    while Q:
        i = Q.pop()
        if i < 0:
            i = ~i
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                k = ET2[i] - ET1[i] + 1 - USED[C[P[i]]] + ORG[i]
                ANS[C[P[i]]] -= f(k)
                TMP[P[i]] += k
            continue
        done[i] = 1
        if i: ORG[i] = USED[C[P[i]]]
        ct += 1
        ET1[i] = ct
        for a in X[i]:
            if done[a]: continue
            P[a] = i
            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
import numpy as np
from numba import njit


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


@njit('i8[:](i8[:],i8[:],i8[:],i8[:],i8[:])', cache=True)
def main(C, tour, ind_L, ind_R, parent):
    N = len(C) - 1
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = np.zeros(N + 1, np.int64)
C[1:] = np.array(readline().split(), np.int64)
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = map(lambda x: np.array(x, np.int64), EulerTour(graph))

answer = main(C, tour, ind_L, ind_R, parent)
print('\n'.join(answer[1:].astype(str)))

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np
from numba import njit


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = np.zeros(N + 1, np.int64)
C[1:] = np.array(readline().split(), np.int64)
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = map(lambda x: np.array(x, np.int64), EulerTour(graph))


@njit('i8[:](i8[:],i8[:],i8[:],i8[:],i8[:])', cache=True)
def main(C, tour, ind_L, ind_R, parent):
    N = len(C) - 1
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


answer = main(C, tour, ind_L, ind_R, parent)
print('\n'.join(answer[1:].astype(str)))

Solution (Python):
import sys
from collections import defaultdict as dd
input = sys.stdin.readline
N = int(input())
c = [0] + list(map(int, input().split()))
e = dd(list)
for _ in range(N - 1):
  u, v = map(int, input().split())
  e[u].append(v)
  e[v].append(u)
s = [1]
order = []
vis = [0] * (N + 1)
ps = [0] * (N + 1)
while len(s):
  x = s.pop()
  vis[x] = 1
  order.append(x)
  for y in e[x]:
    if vis[y]: continue
    s.append(y)
    ps[y] = x

size = [1] * (N + 1)
res = [0] * (N + 1)
d = [dd(int) for _ in range(N + 1)]
for x in order[: : -1]:
  for y in e[x]:
    if y == ps[x]: continue
    size[x] += size[y]
    res[c[x]] += (size[y] - d[y][c[x]]) * (size[y] - d[y][c[x]] + 1) // 2

    if len(d[x]) < len(d[y]): d[x], d[y] = d[y], d[x]
    for z in d[y]: d[x][z] += d[y][z]
  d[x][c[x]] = size[x]


for i in range(1, N + 1):
  print(N * (N + 1) // 2 - res[i] - (N - d[1][i]) * (N - d[1][i] + 1) // 2)

Solution (Python):
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np
from numba import njit


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = np.zeros(N + 1, np.int64)
C[1:] = np.array(readline().split(), np.int64)
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = map(lambda x: np.array(x, np.int64), EulerTour(graph))


@njit(cache=True)
def main(C, tour, ind_L, ind_R, parent):
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


answer = main(C, tour, ind_L, ind_R, parent)
print('\n'.join(answer[1:].astype(str)))

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


def EulerTour(graph, root=1):
    global ind_L, ind_R, parent
    N = len(graph) - 1
    stack = [-root, root]
    tour = []
    i = -1
    while stack:
        v = stack.pop()
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
        yield v
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
parent = [0] * (N + 1)
ind_L = [0] * (N + 1)
ind_R = [0] * (N + 1)
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in EulerTour(graph):
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        c = C[parent[v]]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  
        x -= removed[c] - memo[v]
        answer[c] -= x * (x + 1) // 2
        removed[c] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))

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


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = EulerTour(graph)
removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in tour:
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        c = C[parent[v]]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  
        x -= removed[c] - memo[v]
        answer[c] -= x * (x + 1) // 2
        removed[c] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))

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


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = EulerTour(graph)
removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in tour:
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        p = parent[v]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  
        x -= removed[C[p]] - memo[v]
        answer[C[p]] -= x * (x + 1) // 2
        removed[C[p]] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))