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
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)

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

res = res[::-1]
for v in res:
    for nv in edge[v]:
        if depth[nv]==depth[v]+1:
            size[v] += size[nv]

vertex = []

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

euler_tour(1,0)

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

visit = [-1 for i in range(N+1)]
for v in vertex:
    color = c[v-1]
    if visit[v] == -1:
        if parent[color]!=0:
            color_tree_edge[parent[color]].append(v)
        else:
            color_tree_edge[0][color].append(v)

        visit[v] = parent[color]
        parent[color] = v
    else:
        parent[color] = visit[v]

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

    return res[::-1]

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

for color in range(1,N+1):
    res = N*(N+1)//2

    vertex = bfs(color)

    for v in vertex:
        tmp = size[v]
        if v!=0 and color_tree_edge[v]:
            tmp -= 1
            for nv in color_tree_edge[v]:
                tmp -= size[nv]
            res -= tmp*(tmp+1)//2
        elif v==0:
            for nv in color_tree_edge[v][color]:
                tmp -= size[nv]
            res -= tmp*(tmp+1)//2
        else:
            for nv in edge[v]:
                if depth[nv] == depth[v] + 1:
                    res -= size[nv]*(size[nv]+1)//2

    ans[color] = res

for i in range(1,N+1):
    print(ans[i])

Submission status: Wrong Answer

Solution code:
# 入力のみ
# 入力をいい感じに

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):
    # 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):
    return np.array([1])
    
    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()

Submission status: Wrong Answer

Solution code:
# どこが重いんだこれ
# 入出力だけ

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):
    # 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):
    return np.array([1])
    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())))

Submission status: Wrong Answer

Solution code:
# どこが重いんだこれ
# グラフ構築で切る

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):
    # 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)
    
    return np.array([1])


    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())))

Submission status: Wrong Answer

Solution code:
# どこが重いんだこれ
# オイラーツアーだけで切る

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):
    # 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)
    
    return np.array([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())))

Submission status: Wrong Answer

Solution code:
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]:
        i = c[et[i]] - 1
        val = bit.get_val(i)
        ans[i] += (val // 2) * (val // 2 + 1) // 2

        bit.add(0, n, 1)
        bit.add(i, i + 1, -(val + 1))
        if len(stack[i]) >= 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 // 2) * (val // 2 + 1) // 2

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

Submission status: Wrong Answer

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

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

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

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

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

Submission status: Wrong Answer

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

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

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

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

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

Submission status: Wrong Answer

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

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

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

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

#print(sz)

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

Submission status: Wrong Answer

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
        if not self.children:
            color_stack[self.color].append(self.child_count)
        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 = filter(lambda x:x>0,color_groups[color] + color_stack[color])
    # print(v)
    print(N + N * (N - 1) // 2 - sum(map(lambda x: x + x * (x - 1) // 2, v)))
    # print(v, N + N * (N - 1) // 2)

Submission status: Wrong Answer

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
        if not self.children:
            color_stack[self.color].append(self.child_count)
        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)