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



Sample input:
3
1 2 1
1 2
2 3


Sample output:
5
4
0

Submission status: Accepted

Solution code:
import sys
sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))
sys.setrecursionlimit(10**6)
class Bit:
    def __init__(self, n):
        self.size = n
        self.tree = [0] * (n + 1)
 
    def sum(self, i):
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & (-i)
        return s
 
    def add(self, i, x):
        while i <= self.size:
            self.tree[i] += x
            i += i & (-i)
def resolve():
    N = ir()
    c = lr()
    ab = [[] for i in range(N)]
    for i in range(N-1):
        a, b = lr()
        ab[a-1].append(b-1)
        ab[b-1].append(a-1)
    tot = N*(N+1)//2
    il = [0]*N
    ol = [0]*N
    cnt = [0]
    def dfs(u, p):
        cnt[0] += 1
        il[u] = cnt[0]
        for i in ab[u]:
            if i != p:
                dfs(i, u)
        ol[u] = cnt[0]
    dfs(0, -1)
    vs = [[] for i in range(N)]
    for i in range(N):
        vs[c[i]-1].append(i)
    sl = Bit(N)
    for i in range(1, N+1):
        sl.add(i, 1)
    for i in range(N):
        his = []
        ans = tot
        tvs = sorted(vs[i], key=lambda x: il[x])[::-1]
        for v in tvs:
            n = 1
            for j in ab[v]:
                if il[j] > il[v]:
                    tmp = sl.sum(ol[j])-sl.sum(il[j]-1)
                    ans -= tmp*(tmp+1)//2
                    n += tmp
            sl.add(il[v], -n)
            his.append([il[v], n])
        tmp = sl.sum(ol[0])-sl.sum(0)
        ans -= tmp*(tmp+1)//2
        for h in his:
            sl.add(h[0], h[1])
        print(ans)
resolve()

Submission status: Accepted

Solution code:
import sys
sys.setrecursionlimit(10**6)
def input():
	return sys.stdin.readline()[:-1]

n = int(input())
c = list(map(int, input().split()))
ans = [n*(n+1)//2 for _ in range(n)]
for i in range(n):
	c[i] -= 1
adj = [[] for _ in range(n)]
for _ in range(n-1):
	a, b = map(int, input().split())
	adj[a-1].append(b-1)
	adj[b-1].append(a-1)

s = 0
dp = [0 for _ in range(n)]

def dfs(x, p):
	global s
	pres = s + dp[c[x]]
	for v in adj[x]:
		if v == p:
			continue
		dp[c[x]] = -s
		dfs(v, x)
		ans[c[x]] -= (s+dp[c[x]]) * (s+dp[c[x]]+1) // 2
	s += 1
	dp[c[x]] = pres - s
	return

dfs(0, -1)
for i in range(n):
	ans[i] -= (s+dp[i]) * (s+dp[i]+1) // 2

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

Submission status: Accepted

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

under_col = [0] * N
ans = [0] * N


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


def go(curr, fro):
    col = Cs[curr] - 1
    prev_under = under_col[col]
    deg = 1
    for n in adj[curr]:
        if n == fro:
            continue
        curr_under = under_col[col]
        ch_deg = yield go(n, curr)
        diff_under = under_col[col] - curr_under
        ans[col] += tri(ch_deg - diff_under)
        deg += ch_deg
    under_col[col] = prev_under + deg
    return deg


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


tramp(go(0, -1))
for i in range(N):
    ans[i] += tri(N - under_col[i])
    print(tri(N) - ans[i])

Submission status: Accepted

Solution code:
from collections import defaultdict
import sys

sys.setrecursionlimit(10 ** 6)
int1 = lambda x: int(x) - 1
p2D = lambda x: print(*x, sep="\n")
def II(): return int(sys.stdin.readline())
def MI(): return map(int, sys.stdin.readline().split())
def MI1(): return map(int1, sys.stdin.readline().split())
def LI(): return list(map(int, sys.stdin.readline().split()))
def LI1(): return list(map(int1, sys.stdin.readline().split()))
def LLI(rows_number): return [LI() for _ in range(rows_number)]
def SI(): return sys.stdin.readline()[:-1]

class BitSum:
    def __init__(self, n):
        self.n = n + 3
        self.table = [0] * (self.n + 1)

    def add(self, i, x):
        i += 1
        while i <= self.n:
            self.table[i] += x
            i += i & -i

    def sum(self, i):
        i += 1
        res = 0
        while i > 0:
            res += self.table[i]
            i -= i & -i
        return res

def main():
    def dfs(u=0,pu=-1):
        ll[u]=len(uu)
        uu.append(u)
        for v in to[u]:
            if v==pu:continue
            dfs(v,u)
        rr[u]=len(uu)

    def sumlr(l,r):
        return bit.sum(r-1)-bit.sum(l-1)

    n=II()
    cc=LI1()
    to=[[] for _ in range(n)]
    for _ in range(n-1):
        a,b=MI1()
        to[a].append(b)
        to[b].append(a)
    ll=[-1]*n
    rr=[-1]*n
    uu=[]
    dfs()
    bit=BitSum(n)
    for u in range(n):bit.add(u,1)
    ctou=defaultdict(list)
    for u,c in enumerate(cc):ctou[c].append(u)
    pn=lambda x:x*(x+1)//2
    for c in range(n):
        log=[]
        ans = pn(n)
        for u in sorted(ctou[c],key=lambda x:-ll[x]):
            us=1
            for v in to[u]:
                if ll[v]<ll[u]:continue
                vs=sumlr(ll[v],rr[v])
                ans-=pn(vs)
                us+=vs
            bit.add(ll[u],-us)
            log.append((ll[u],us))
        vs = sumlr(0,n)
        ans -= pn(vs)
        print(ans)
        for i,s in log:bit.add(i,s)

main()

Submission status: Accepted

Solution code:
import math
import random
import sys
from bisect import bisect_left, bisect_right
from collections import Counter, defaultdict, deque
from functools import cmp_to_key, lru_cache, reduce
from operator import add, sub
from itertools import count

sys.setrecursionlimit(10**6 + 1)
input = sys.stdin.buffer.readline

MOD = 10**9 + 7


def solve(g, clr, n):
    def dfs(v, par):
        nonlocal total
        sub[v] = 1
        bfr = cnt[clr[v]]

        for u in g[v]:
            if u == par:
                continue
            pre = total - cnt[clr[v]]
            sub[v] += dfs(u, v)
            nxt = total - cnt[clr[v]]
            ans[clr[v]] -= (nxt-pre) * (nxt-pre-1) // 2

        cnt[clr[v]] = bfr + sub[v]
        total += 1
        return sub[v]

    ans = [(n-1)*n//2] * (n+1)
    cnt = [0] * (n+1)
    sub = [0] * (n+1)
    total = 0

    dfs(1, -1)
    # print(cnt)
    for i in range(1, n+1):
        pre, nxt = 0, n-cnt[i]
        ans[i] -= (nxt-pre) * (nxt-pre-1) // 2

    for c in clr:
        ans[c] += 1
    print(*ans[1:], sep='\n')


if __name__ == "__main__":
    n = int(input())
    colors = [0] + [int(s) for s in input().split()]
    g = defaultdict(set)

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

    solve(g, colors, n)

Submission status: Accepted

Solution code:
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)]
head = [-1] * (N + 1)
to = [0] * (N - 1 << 1)
nxt = [0] * (N - 1 << 1)
for i in range(N-1):
    x, y = map(int, input().split())
    x, y = x-1, y-1
    nxt[i] = head[x]
    to[i] = y
    head[x] = i
    j = i + N - 1
    nxt[j] = head[y]
    to[j] = x
    head[y] = j

def EulerTour(n, i=0):
    f = lambda k: k * (k + 1) // 2
    USED = [0] * n
    ORG = [0] * n
    TMP = [0] * n
    ANS = [f(n)] * n
    
    P = [-1] * n
    ct = 0
    ET1 = [0] * n
    ET2 = [0] * n
    while i >= 0:
        e = head[i]
        if e < 0:
            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 = to[e]
        if P[i] == j:
            head[i] = nxt[e]
            continue
        
        P[j] = i
        head[i] = nxt[e]
        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, 0), sep = "\n")

Submission status: Accepted

Solution code:
import sys
stdin = sys.stdin
sys.setrecursionlimit(10**6)
ni = lambda: int(ns())
na = lambda: list(map(int, stdin.readline().split()))
nn = lambda: list(stdin.readline().split())
ns = lambda: stdin.readline().rstrip()

n = ni()
c = [i-1 for i in na()]
g = [[] for i in range(n)]

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

def EulerTour(n, X, i0):
    # Xは破壊してXとPができる
    P = [-1] * n
    Q = [~i0, i0]
    ct = -1
    x = [0]*n
    y = [0]*n
    z = [0]*n
    ET = []
    ET1 = [0] * n
    ET2 = [0] * n
    DE = [0] * n
    de = -1
    ans = [n*(n+1)//2]*n
    while Q:
        i = Q.pop()
        if i < 0:
            i = ~i
            ET2[i] = ct
            de -= 1
            x[c[i]] += 1 + z[i]
            if i:
                p = P[i]
                k = ET2[i] - ET1[i] + 1 - x[c[p]] + y[i]
                ans[c[p]] -= k*(k+1)//2
                z[p] += k
            continue
        if i >= 0:
            ET.append(i)
            ct += 1
            if i:
                y[i] = x[c[P[i]]]
            if ET1[i] == 0: ET1[i] = ct
            de += 1
            DE[i] = de
        for a in X[i][::-1]:
            if a != P[i]:
                P[a] = i
                for k in range(len(X[a])):
                    if X[a][k] == i:
                        del X[a][k]
                        break
                Q.append(~a)
                Q.append(a)
    for i in range(n):
        ans[i] -= (n - x[i])*(n - x[i]+1)//2
    return ans

print(*EulerTour(n,g,0), sep='\n')

Submission status: Accepted

Solution code:
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")

Submission status: Accepted

Solution code:
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")

Submission status: Accepted

Solution code:
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")

Submission status: Accepted

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


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

# ci = [set() for _ in range(n)]
# for i,c in enumerate(C):
#     ci[c-1].add(i)

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
# index:color
reached = [False]*n
size = [0]*n
# index:vertex

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()
    # size[p] += s
    if cparent[c]:
        size[cparent[c][-1]] -= s
    else:
        root_size[c] -= s
    # ans[c] -= size[p] * (size[p] - 1)//2
    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)

Submission status: Accepted

Solution code:
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  # subtree size
        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:])))

Submission status: Accepted

Solution code:
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  # subtree size
            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:])))

Submission status: Accepted

Solution code:
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  # subtree size
            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)))

Submission status: Accepted

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

Submission status: Accepted

Solution code:
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  # subtree size
            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)))

Submission status: Accepted

Solution code:
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")

Submission status: Accepted

Solution code:
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  # subtree size
            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)))