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

sys.setrecursionlimit(2 * 10 ** 5)

(N,) = [int(x) for x in input().split()]
colors = [int(x) for x in input().split()]
edges = [[int(x) for x in input().split()] for i in range(N - 1)]
graph = defaultdict(list)
for u, v in edges:
    graph[u].append(v)
    graph[v].append(u)
parent = {1: None}
children = defaultdict(list)


def rootTree(node):
    for nbr in graph[node]:
        if nbr not in parent:
            parent[nbr] = node
            children[node].append(nbr)
            rootTree(nbr)


rootTree(1)


colorlessCount = [0 for c in range(N + 1)]
colorlessTime = [0 for c in range(N + 1)]
t = 0


def buildComps(node):
    global t
    t += 1
    currColor = colors[node - 1]
    colorlessTime[currColor] += 1
    for child in children[node]:
        oldCount = t - colorlessTime[currColor]
        colorlessTime[currColor] = t
        buildComps(child)
        s = t - colorlessTime[currColor]
        colorlessTime[currColor] = t - oldCount
        colorlessCount[currColor] += s * (s + 1) // 2


buildComps(1)

for c in range(1, N + 1):
    s = t - colorlessTime[c]
    colorlessCount[c] += s * (s + 1) // 2
maxPairs = N * (N + 1) // 2
for c in range(1, N + 1):
    print(maxPairs - colorlessCount[c])

Submission status: Accepted

Solution code:
import os
from collections import Counter, defaultdict

import sys

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353

N = int(sys.stdin.buffer.readline())
C = list(map(int, sys.stdin.buffer.readline().split()))
AB = [list(map(int, sys.stdin.buffer.readline().split())) for _ in range(N - 1)]

graph = [[] for _ in range(N)]
for a, b in AB:
    a -= 1
    b -= 1
    graph[a].append(b)
    graph[b].append(a)

color_counts = [defaultdict(int) for _ in range(N + 1)]
for c in range(1, N + 1):
    color_counts[c][0] = N
color_v_weights = [[] for _ in range(N + 1)]
color_v = [0 for _ in range(N + 1)]
color_v_next = [1 for _ in range(N + 1)]


# color_parents = [[] for _ in range(N + 1)]
# color_stack = [[] for _ in range(N + 1)]
# children_counts = defaultdict(int)


def dfs(v, p):
    color = C[v]
    color_v_prev = color_v[color]

    children_cv = []
    ret = 1
    for u in graph[v]:
        if u == p:
            continue
        cv = color_v_next[color]
        children_cv.append(cv)
        color_v[color] = cv
        color_v_next[color] += 1
        cnt = dfs(u, v)
        ret += cnt
        color_counts[color][cv] += cnt
    color_v[color] = color_v_prev
    color_counts[color][color_v_prev] -= ret
    return ret


dfs(0, None)
counter = Counter(C)

for color, cc in enumerate(color_counts[1:], start=1):
    ans = N * (N - 1) // 2 + counter[color]
    for cnt in cc.values():
        ans -= cnt * (cnt - 1) // 2
    print(ans)

Submission status: Accepted

Solution code:
import sys
sys.setrecursionlimit(10**9)

mod = 1000000007
INF = 110110110

# input
N = int(input())
c = list(map(int, input().split()))
for i in range(N):
    c[i] -= 1

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


# solve
child = [0] * N

def dfs(now, par):
    res = 1
    for to in E[now]:
        if to == par:
            continue
        res += dfs(to, now)
    child[now] = res
    return res

dfs(0, -1)

ans = [0] * N
ok = [N] * N
def solve(now, par):
    memo = ok[c[now]]
    up = ok[c[now]] - child[now]
    num = [up]
    for to in E[now]:
        if to == par:
            continue
        num.append(child[to])
    S = sum(num)
    ans[c[now]] += S + 1
    for n in num:
        S -= n
        ans[c[now]] += S * n

    for to in E[now]:
        if to == par:
            continue
        ok[c[now]] = child[to]
        solve(to, now)
    ok[c[now]] = memo - child[now]

solve(0, -1)

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

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.readline

n = int(input())
color = list(map(lambda x: int(x)-1,input().split()))
out = [0] * n

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

stacks = [[[0,0]] for i in range(n)]
segs = [[] for i in range(n)]

adj = [[] for i in range(n)]
visited = [False] * n
dfs = []

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

visited[0] = True
dfs.append(0)
stacks[color[0]].append([1,0])
curr = 1

while dfs:
    node = dfs[-1]
    prev, rem = stacks[color[node]][-1]
    diff = (curr - prev - rem)
    out[color[node]] += (diff * (diff + 1))//2
    stacks[color[node]][-1][1] += (diff) 
    
    if adj[node]:
        nex = adj[node].pop()
        if not visited[nex]:
            visited[nex] = True
            dfs.append(nex)
            curr += 1
            stacks[color[nex]].append([curr,0])
    else:
        dfs.pop()
        prev, rem = stacks[color[node]].pop()
        stacks[color[node]][-1][1] += rem + 1
for i in range(n):
    diff = (curr - stacks[i][0][1])
    out[i] += (diff * (diff + 1))//2
    out[i] = total - out[i]

print('\n'.join(map(str,out)))

Submission status: Runtime Error

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 sort_edges(edges_):
    M = len(edges_)
    edges = np.empty((M * 2, 2), dtype=np.int64)
    edges[:M] = edges_
    edges[M:] = edges_[:, ::-1]
    order = np.argsort(edges[:, 0])  # O(N) にできなくもない
    edges = edges[order, 1]
    c = np.cumsum(np.bincount(edges_.ravel()))
    lefts = np.zeros(len(c) + 1, dtype=np.int64)
    lefts[1:] = c
    return edges, lefts

def eular_tour(edges, lefts, root):
    # グラフは 1-indexed が良い
    n = len(lefts)-1
    stack = [root]
    tour = [0] * 0
    firsts = np.full(n, -100, dtype=np.int64)
    lasts = np.full(n, -100, dtype=np.int64)
    parents = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if firsts[v] >= 0:
            lasts[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        p = parents[v]
        firsts[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in edges[lefts[v]:lefts[v+1]]:
            if p != u:
                parents[u] = v
                stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, firsts, lasts, parents

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

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

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

    return Ans[1:]


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


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

main()

Submission status: Runtime Error

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 = [[0] 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))
        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: Runtime Error

Solution code:
def euler_tour(tree: list, root: int):
    """オイラーツアーを行う"""
    n = len(tree)
    par = {root: -1}
    trail = []
    stack = [root]
    while stack:
        v = stack.pop()
        if v >= 0:  # 行きがけ順の処理
            trail.append(v)
            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])
 
    return trail, par


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, par = euler_tour(tree, root)

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

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

    if par[et[i - 1]] == et[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))
        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: Runtime Error

Solution code:
import numpy as np

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: Runtime Error

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) 

counter = 0
ds = [0 for _ in range(N)]
rs = [0 for _ in range(N)]

# 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: Runtime Error

Solution code:

def s0():return input()
def s1():return input().split()
def s2(n):return [input() for x in range(n)]
def s3(n):return [[input().split()] for _ in range(n)]
def n0():return int(input())
def n1():return [int(x) for x in input().split()]
def n2(n):return [int(input()) for _ in range(n)]
def n3(n):return [[int(x) for x in input().split()] for _ in range(n)]

n=n0()
c=n1()
ab=[(x[0],x[1]) for x in n3(n-1)]
 
import networkx as nx
G = nx.Graph()
 
G.add_nodes_from(range(1,n+1))
G.add_edges_from(ab)
 
dict_color={i:0 for i in range(1,n+1)}
for j in c:
    dict_color[j]+=1

for i in range(1,n+1):
    for line in list(nx.all_simple_paths(G, source=i, target=[x for x in range(i+1,n+1)])):
        path_color = [c[x-1] for x in line]
        for i in set(path_color):
            dict_color[i]+=1

for i in dict_color:
    print(dict_color[i])

Submission status: Runtime Error

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)

Submission status: Runtime Error

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]
    if any(x < 0 for x in v):
        raise
    # 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: Runtime Error

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:
            if trees[node_index].color == self.color:
                raise
            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(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: Runtime Error

Solution code:
import bisect
import collections
import sys

sys.setrecursionlimit(100000)
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)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(2*10**5)

N = int(input())  # 頂点数
C = [int(x) - 1  for x in input().split()]  # 頂点の色番号 (1~N) -> (0~N-1)

E = [[] for _ in range(N)]
for _ in range(N-1):
    a, b = map(int, input().split())
    E[a-1].append(b-1)
    E[b-1].append(a-1)
    
def dfs(u, p = -1):

    color = C[u]  # 色番号
    VC[color].append(u)  # 色ごとに stack しておく
    
    s = 1  # 部分木サイズ（を格納する）
    for v in E[u]:
        if v == p:
            continue
            
        child[u] = 0  # v の方向にある子孫のうち辿れる頂点の数 = s - 子孫にある、同色を根とする部分木のサイズの合計
        ret = dfs(v, u)

        s += ret  # 純粋な部分木のサイズを計算
        child[u] += ret  # 同色の頂点が切り離されたあとの部分木サイズ
        
        ans[color] -= child[u] * (child[u] + 1) // 2        

    VC[color].pop()
    if VC[color]: # 上位に同じ色がある場合
        child[VC[color][-1]] -= s  # 一番近い同色から、自分の部分木の大きさをひく
    else:  # 上位に同じ色がない場合
        root_size[color] -= s  # 上位に残っている頂点の数
        
    return s

child = [0] * N  # 辿っている子の中でのの残り数　とりあえず 0で初期化
root_size = [N] * N  # 根方向に残っている頂点の数（を格納する。操作前は N 全部）
VC = [ [] for _ in range(N)]  # 色番号ごと
ans = [N*(N+1)//2] * N  # 色ごとの（残りの）組み合わせ数　初期値は 全点の組み合わせ

dfs(0)

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

Submission status: Runtime Error

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
    for i in range(1, n + 1):
        ans = solve(graph, i, stack)
        print(ans)

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