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: Wrong Answer

Solution code:
import sys

stdin = sys.stdin
sys.setrecursionlimit(10**6)
ns = lambda: stdin.readline().rstrip()
ni = lambda: int(stdin.readline().rstrip())
nm = lambda: map(int, stdin.readline().split())
nl = lambda: list(map(int, stdin.readline().split()))

n = ni()
color = nl()
g = [list() for _ in range(n)]
for _ in range(n-1):
    a, b = nm()
    g[a-1].append(b-1)
    g[b-1].append(a-1)

size = [1]*n
h = [list() for _ in range(n+1)]
sst = [0]*n
def dfs(v, p):
    h[color[v]].append(v)
    for x in g[v]:
        if x == p: continue
        size[v] += dfs(x, v)
    h[color[v]].append(v)
    return size[v]
dfs(0, -1)
# print(size)
acc = [0]*(n+10)
ans = [n*(n+1)//2]*(n+1)
for col in range(1, n+1):
    # print(h[col])
    acc[-1] = 0
    p = [-1]
    for v in h[col]:
        if p[-1] == v:
            p.pop()
            q = size[v] - acc[v] - 1
            ans[col] -= q*(q+1)//2
            # print(v, q)
            acc[p[-1]] += size[v]
        else:
            p.append(v)
    q = n - acc[-1]
    # print(-1, q)
    ans[col] -= q*(q+1)//2
print(*ans[1:], sep='\n')

Submission status: Time Limit Exceeded

Solution code:
# 全部非再帰に書き換え・numbaなし

import numpy as np

def eular_tour(G, root):
    # G は破壊的に有向木にされる
    # グラフは 1-indexed が良い
    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)))

Submission status: Time Limit Exceeded

Solution code:
import sys


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


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

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

    def each_child(parent, root, child):
        if parent == child:
            return
        dfs(root, child)
        size[root] += size[child]

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

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

        list(map(lambda x: each_child(parent, root, x), adj[root]))

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

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

        out_time[root] = timer
        timer += 1

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

    sys.setrecursionlimit(10**6)

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


solve()

Submission status: Time Limit Exceeded

Solution code:
import sys


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


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

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

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

        def each_child(child):
            if parent == child:
                return
            dfs(root, child)
            size[root] += size[child]

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

        list(map(each_child, adj[root]))

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

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

        out_time[root] = timer
        timer += 1

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

    sys.setrecursionlimit(10**6)

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


solve()

Submission status: Time Limit Exceeded

Solution code:
import sys
import numpy as np
sys.setrecursionlimit(10**6)

N = int(input())
Vs = list(map(int, input().split()))

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

counter = 0
ds = np.zeros((N,), dtype=np.int32)
rs = np.zeros((N,), dtype=np.int32)

# euler tour
def step(c, p):
    global counter
    ns = Gs[c-1]
    v = Vs[c-1] 

    counter += 1

    for n in ns:
        if n == p:
            continue
        cc = counter
        pd = ds[v-1]

        step(n, c)

        d = counter - cc
        w = d - (ds[v-1] - pd)

        ds[v-1] = pd + d
        rs[v-1] += int(w * (w-1) / 2 + w)

    ds[v-1] += 1

step(1, -1)

t = int(N * (N-1) / 2 + N)
for k in range(N):
    d = ds[k]
    w = counter - d
    rs[k] += int(w * (w-1) / 2 + w)
    print(t - rs[k])

Submission status: Time Limit Exceeded

Solution code:
N = int(input())
Vs = list(map(int, input().split()))

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

Ks = [[0] for _ in range(N)]

# euler tour
def step(c, p):
    ns = Gs[c-1]
    v = Vs[c-1] 

    for k in range(N):
        if k != v-1:
            Ks[k][-1] += 1

    for n in ns:
        if n == p:
            continue
        Ks[v-1].append(0)
        step(n, c)

step(1, -1)

t = int(N * (N-1) / 2 + N)
for k in range(N):
    ss = Ks[k]
    counter = 0
    for s in ss:
        counter += s * (s-1) / 2 + s
    print(t - int(counter))

Submission status: Time Limit Exceeded

Solution code:
import sys
input = sys.stdin.readline
# 文字列をinput()した場合、末尾に改行が入るので注意
def main():
    n=int(input())
    c_=list(map(int,input().split()))
    ab=[list(map(int,input().split())) for _ in range(n-1)]
    """
    n=10000
    c_=list(range(1,n+1))
    ab=[[i+1,i+2] for i in range(n-1)]
    """
    c=[[] for _ in range(n)]
    for i,ci in enumerate(c_):
        c[ci-1].append(i)

    ki=[[] for _ in range(n)]
    for a,b in ab:
        ki[a-1].append(b-1)
        ki[b-1].append(a-1)
    base=n+(n*(n-1))//2
    def nopass(nol):
        if len(nol)==0:
            return 0
        elif len(nol)==n:
            return base
        start=set()
        rn=1 #連結成分の個数
        seen=[0]*n
        for li in nol:
            seen[li]=1
            start |= set(ki[li])
            rn+=len(ki[li])-1
        start-=set(nol)
        ret=0
        cn=0
        nn=len(nol)
        nrn=0
        seenp=len(nol)
        for i in range(n):
            if seen[i]==1:
                continue
            c=0
            nrn+=1
            if nrn==rn:
                c=n-seenp
                ret+=c+(c*(c-1))//2
                break
            todo=[i]
            while todo:
                a=todo.pop()
                seen[a]=1
                seenp+=1
                c+=1
                l=ki[a]
                for li in l:
                    if seen[li]==0:
                        todo.append(li)
            ret+=c+(c*(c-1))//2
        return base-ret
    for ci in c:
        print(nopass(ci))

if __name__=='__main__':
    main()

Submission status: Time Limit Exceeded

Solution code:
n=int(input())
c_=list(map(int,input().split()))
ab=[list(map(int,input().split())) for _ in range(n-1)]
"""
n=10000
c_=list(range(1,n+1))
ab=[[i+1,i+2] for i in range(n-1)]
"""
c=[[] for _ in range(n)]
for i,ci in enumerate(c_):
    c[ci-1].append(i)

ki=[[] for _ in range(n)]
for a,b in ab:
    ki[a-1].append(b-1)
    ki[b-1].append(a-1)
base=n+(n*(n-1))//2
def nopass(nol):
    if len(nol)==0:
        return 0
    elif len(nol)==n:
        return base
    seen=[0]*n
    for li in nol:
        seen[li]=1
    ret=0
    for i in range(n):
        if seen[i]==1:
            continue
        c=0
        todo=[i]
        while todo:
            a=todo.pop()
            seen[a]=1
            c+=1
            l=ki[a]
            for li in l:
                if seen[li]==0:
                    todo.append(li)
        ret+=c+(c*(c-1))//2
    return base-ret
for ci in c:
    print(nopass(ci))

Submission status: Time Limit Exceeded

Solution code:
n=int(input())
c=list(map(int,input().split()))
ab=[list(map(int,input().split())) for _ in range(n-1)]
ki=[[] for _ in range(n)]
for a,b in ab:
    ki[a-1].append(b-1)
    ki[b-1].append(a-1)
base=n+(n*(n-1))//2
def nopass(nol):
    if len(nol)==0:
        return base
    elif len(nol)==n:
        return 0
    seen=[0]*n
    for li in nol:
        seen[li]=1
    ret=0
    for i in range(n):
        if seen[i]==1:
            continue
        c=0
        todo=[i]
        while todo:
            a=todo.pop()
            seen[a]=1
            c+=1
            l=ki[a]
            for li in l:
                if seen[li]==0:
                    todo.append(li)
        ret+=c+(c*(c-1))//2
    return ret
for i in range(n):
    nol=[j for j,cj in enumerate(c) if cj==i+1]
    print(base-nopass(nol))

Submission status: Time Limit Exceeded

Solution code:
import bisect
import collections
import sys

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


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


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


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

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

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

visited = [False] * N

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

# dfs = [0]
# while dfs:
#     node = dfs[-1]
#     prev, rem = stacks[c[node]][-1]
#     diff = curr - prev - rem
#     ans[c[node]] += (diff * (diff + 1)) // 2
#     stacks[c[node]][-1][1] += diff
#     if adj[node]:
#         next_node = adj[node].pop()
#         if not visited[next_node]:
#             visited[next_node] = True


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


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

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

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


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

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

Submission status: Time Limit Exceeded

Solution code:
import bisect
import collections
import sys

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


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


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


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

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

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

visited = [False] * N

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

# dfs = [0]
# while dfs:
#     node = dfs[-1]
#     prev, rem = stacks[c[node]][-1]
#     diff = curr - prev - rem
#     ans[c[node]] += (diff * (diff + 1)) // 2
#     stacks[c[node]][-1][1] += diff
#     if adj[node]:
#         next_node = adj[node].pop()
#         if not visited[next_node]:
#             visited[next_node] = True


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


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

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

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


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

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

Submission status: Time Limit Exceeded

Solution code:
import bisect
import collections
import sys

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


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


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


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

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

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

visited = [False] * N

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

# dfs = [0]
# while dfs:
#     node = dfs[-1]
#     prev, rem = stacks[c[node]][-1]
#     diff = curr - prev - rem
#     ans[c[node]] += (diff * (diff + 1)) // 2
#     stacks[c[node]][-1][1] += diff
#     if adj[node]:
#         next_node = adj[node].pop()
#         if not visited[next_node]:
#             visited[next_node] = True


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


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

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

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


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

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

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
sys.setrecursionlimit(10**7)
from pprint import pprint as pp
from pprint import pformat as pf
# @pysnooper.snoop()
#import pysnooper # debug

import math
#from sortedcontainers import SortedList, SortedDict, SortedSet # no in atcoder
import bisect

class Graph:

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

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

    def get_connect(self, v):
        return self.edges[v]

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



def solve(graph, color, stack):
    ans = calc(graph.size)
    stack_key = 0
    stack[stack_key] = 0
    #print('stack') # debug
    #print(stack) # debug
    def dfs(v, parent):
        nonlocal ans, stack, stack_key
        #print('v', v) # debug
        if graph.vertices[v] == color:
            stack_key += 1
            stack[stack_key] = 0
        else:
            stack[stack_key] += 1
        for to in graph.edges[v]:
            if to == parent:
                continue
            dfs(to, v)
            if graph.vertices[v] == color:
                ans -= calc(stack[stack_key])
                stack[stack_key] = 0
        if graph.vertices[v] == color:
            stack_key -= 1
    dfs(0, -1)
    ans -= calc(stack[stack_key])
    return ans





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

    stack = [None] * (graph.size + 1)
    for i in range(1, n + 1):
        ans = solve(graph, i, stack)
        print(ans)

    #print('\33[32m' + 'end' + '\033[0m') # debug

Submission status: Time Limit Exceeded

Solution code:
#!/usr/bin/env python
# -*- coding: utf-8 -*-

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

import math
#from sortedcontainers import SortedList, SortedDict, SortedSet # no in atcoder
import bisect

class BIT: # binary indexed tree
    """
    http://hos.ac/slides/20140319_bit.pdf
    """

    def __init__(self, size):
        self.size = size
        self.data = [0] * (size + 1) # var[0] is dummy

    def add(self, pos, val):
        assert(pos > 0)
        k = pos
        while k <= self.size:
            self.data[k] += val
            # for next
            k += k & -k

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

    def make_copy(self):
        bit = BIT(self.size)
        bit.data = self.data[:]
        return bit

class Graph:

    def __init__(self, size):
        self.size = size
        self.vertices = [0] * (size + 1) # var[0] is dummy
        self.edges = [None] * (size + 1) # var[0] is dummy
        for i in range(1, size + 1):
            self.edges[i] = []

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

    def vertex_ids(self):
        return range(1, len(self.vertices))

    def get_connect(self, v):
        return self.edges[v]


class Solver:

    def __init__(self, graph):
        self.graph = graph
        self.bit = BIT(graph.size)
        self.in_s = [None] * (graph.size + 1) # dummy 0 key
        self.out_s = [None] * (graph.size + 1) # dummy 0 key
        self.parent_s = [None] * (graph.size + 1) # dummy 0 key
        self.color_s =  [None] * (graph.size + 1) # dummy 0 key
        for i, _ in enumerate(self.color_s):
            self.color_s[i] = []

    def run(self):
        self.init_by_dfs()
        #print('self.in_s') # debug
        #print(self.in_s) # debug
        #print('self.out_s') # debug
        #print(self.out_s) # debug
        #print('self.parent_s') # debug
        #print(self.parent_s) # debug
        self.init_color_s()
        #print('self.color_s') # debug
        #print(self.color_s) # debug
        base_bit = self.make_base_bit()
        whole = self.calc(self.graph.size)
        #print('whole') # debug
        #print(whole) # debug
        for colored_v_s in self.color_s[1:]:
            #print('colored_v_s') # debug
            #print(colored_v_s) # debug
            bit = base_bit.make_copy()
            complement = self.solve(colored_v_s, bit)
            #print('complement') # debug
            #print(complement) # debug
            #print('ans') # debug
            print(whole - complement)

    def init_by_dfs(self):
        key = 1
        def dfs(v, parent):
            nonlocal key
            self.parent_s[v] = parent
            self.in_s[v] = key
            for to in self.graph.edges[v]:
                if to == parent:
                    continue
                key += 1
                dfs(to, v)
            self.out_s[v] = key
        dfs(1, 0)

    def init_color_s(self):
        for i, _ in enumerate(self.color_s):
            self.color_s[i] = []
        for v in self.in_s[::-1]:
            if v is None:
                continue
            c = self.graph.vertices[v]
            self.color_s[c].append(v)

    def make_base_bit(self):
        bit = BIT(self.graph.size)
        for i in range(1, self.graph.size + 1):
            bit.add(i, 1)
        return bit

    def solve(self, colored_v_s, bit):
        """
        calc path without colored node
        """
        #print('bit') # debug
        #print(bit.data) # debug
        ans = 0
        for v in colored_v_s:
            #print('v') # debug
            #print(v) # debug
            for to in self.graph.edges[v]:
                if to == self.parent_s[v]:
                    continue
                #print('to') # debug
                #print(to) # debug
                connected = self.calc_connected(to, bit)
                #print('connected') # debug
                #print(connected) # debug
                ans += self.calc(connected)
            self.cut(v, bit)
            #print('bit') # debug
            #print(bit.data) # debug
        connected = self.calc_connected(1, bit)
        ans += self.calc(connected)
        return ans

    def calc_connected(self, v, bit):
        large_key = self.out_s[v]
        large = bit.sum(large_key)
        small_key = self.in_s[v] - 1
        small = bit.sum(small_key)
        return large - small

    def calc(self, a):
        return a * (a + 1) // 2

    def cut(self, v, bit):
        connected = self.calc_connected(v, bit)
        key = self.in_s[v]
        bit.add(key, -1 * connected)









if __name__ == '__main__':
    n = int(input())
    graph = Graph(n)
    graph.vertices = [0] + list(map(int, input().split()))
    for i in range(n - 1):
        frm, to = (map(int, input().split()))
        graph.add_edge(frm, to)
    ans = Solver(graph).run()

    #print('\33[32m' + 'end' + '\033[0m') # debug