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:
# オイラーツアーのテスト

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
    first = np.full(n, -100, dtype=np.int64)
    last = np.full(n, -100, dtype=np.int64)
    while stack:
        v = stack.pop()
        if first[v] >= 0:
            last[v] = len(tour)
            tour.append(-v)  # 帰りがけの辺の表現をマイナス以外にしたい場合ここを変える
            continue
        first[v] = len(tour)
        tour.append(v)
        stack.append(v)
        for u in G[v]:
            G[u].remove(v)
            stack.append(u)
    tour = np.array(tour, dtype=np.int64)
    return tour, first, last

def dfs1(v, G, Siz):
    tour, first, last = eular_tour(G, 0)
    Siz[:] = last - first + 1 >> 1
    return 0

# def dfs1(v, G, Siz):
#     siz = 1
#     for u in G[v]:
#         G[u].remove(v)
#         siz += dfs1(u, G, Siz)
#     Siz[v] = siz
#     return siz

def dfs(v, N, C, G, Cnt, Siz, Ans):
    color = C[v]
    cnt_old = Cnt[color]
    siz = 1
    siz_up = N
    for u in G[v]:
        Cnt[color] = N - Siz[u]  # 上を塞ぐ
        s = dfs(u, N, C, G, Cnt, Siz, Ans)
        siz += s
        siz_up -= s
        Ans[color] += s * (siz_up - cnt_old)
    Ans[color] += siz_up - cnt_old
    Cnt[color] = cnt_old + siz
    return siz

def solve(N, C, AB):
    G = [[0]*0 for _ in range(N)]
    for i in range(N-1):
        a, b = AB[i]
        G[a].append(b)
        G[b].append(a)
    Cnt = np.zeros(N, dtype=np.int64)
    Ans = np.zeros(N, dtype=np.int64)
    Siz = np.zeros(N, dtype=np.int64)
    dfs1(0, G, Siz)
    dfs(0, N, C, G, Cnt, Siz, Ans)
    return Ans

numba_compile([
    [eular_tour, "Tuple((i8[:],i8[:],i8[:]))(List(List(i8)),i8)"],
    [dfs1, "i8(i8,List(List(i8)),i8[:])"],
    [dfs, "i8(i8,i8,i8[:],List(List(i8)),i8[:],i8[:],i8[:])"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])

import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = np.array(input().split(), dtype=np.int64) - 1
AB = np.array([list(map(int, input().split())) for _ in range(N-1)], dtype=np.int64) - 1
Ans = solve(N, C, AB)
print("\n".join(map(str, Ans)))

Submission status: Accepted

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 dfs1(v, G, Siz):
    siz = 1
    for u in G[v]:
        G[u].remove(v)
        siz += dfs1(u, G, Siz)
    Siz[v] = siz
    return siz

def dfs(v, N, C, G, Cnt, Siz, Ans):
    color = C[v]
    cnt_old = Cnt[color]
    siz = 1
    siz_up = N
    for u in G[v]:
        Cnt[color] = N - Siz[u]  # 上を塞ぐ
        s = dfs(u, N, C, G, Cnt, Siz, Ans)
        siz += s
        siz_up -= s
        Ans[color] += s * (siz_up - cnt_old)
    Ans[color] += siz_up - cnt_old
    Cnt[color] = cnt_old + siz
    return siz

def solve(N, C, AB):
    G = [[0]*0 for _ in range(N)]
    for i in range(N-1):
        a, b = AB[i]
        G[a].append(b)
        G[b].append(a)
    Cnt = np.zeros(N, dtype=np.int64)
    Ans = np.zeros(N, dtype=np.int64)
    Siz = np.zeros(N, dtype=np.int64)
    dfs1(0, G, Siz)
    dfs(0, N, C, G, Cnt, Siz, Ans)
    return Ans

numba_compile([
    [dfs1, "i8(i8,List(List(i8)),i8[:])"],
    [dfs, "i8(i8,i8,i8[:],List(List(i8)),i8[:],i8[:],i8[:])"],
    [solve, "i8[:](i8,i8[:],i8[:,:])"],
])

import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = np.array(input().split(), dtype=np.int64) - 1
AB = np.array([list(map(int, input().split())) for _ in range(N-1)], dtype=np.int64) - 1
Ans = solve(N, C, AB)
print("\n".join(map(str, Ans)))

Submission status: Accepted

Solution code:
import sys
input = sys.stdin.buffer.readline
sys.setrecursionlimit(505050)
N = int(input())
C = [int(c)-1 for c in input().split()]
G = [[] for _ in range(N)]
for _ in range(N-1):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    G[a].append(b)
    G[b].append(a)

Cnt = [0] * N
Ans = [0] * N
Siz = [0] * N

def dfs1(v):
    siz = 1
    for u in G[v]:
        G[u].remove(v)
        siz += dfs1(u)
    Siz[v] = siz
    return siz

#Cnt_down = [0] * N

def dfs(v):
    color = C[v]
    cnt_old = Cnt[color]
    siz = 1
    siz_up = N
    #cnt_down = 0
    for u in G[v]:
        #Cnt_down[color] = 0
        Cnt[color] = N - Siz[u]  # 上を塞ぐ
        s = dfs(u)
        siz += s
        siz_up -= s
        Ans[color] += s * (siz_up-cnt_old)
        #cnt_down += Cnt_down[color]
    Ans[color] += siz_up - cnt_old
    #print(siz_up, cnt_old, cnt_down)
    Cnt[color] = cnt_old + siz
    #print(v, Ans)
    return siz

dfs1(0)
dfs(0)
print("\n".join(map(str, Ans)))

Submission status: Accepted

Solution code:
import sys
INF = 1 << 60
MOD = 10**9 + 7 # 998244353
sys.setrecursionlimit(2147483647)
input = lambda:sys.stdin.readline().rstrip()

class SegmentTree(object):
    def __init__(self, A, dot, unit):
        n = 1 << (len(A) - 1).bit_length()
        tree = [unit] * (2 * n)
        for i, v in enumerate(A):
            tree[i + n] = v
        for i in range(n - 1, 0, -1):
            tree[i] = dot(tree[i << 1], tree[i << 1 | 1])
        self._n = n
        self._tree = tree
        self._dot = dot
        self._unit = unit

    def __getitem__(self, i):
        return self._tree[i + self._n]

    def update(self, i, v):
        i += self._n
        self._tree[i] = v
        while i != 1:
            i >>= 1
            self._tree[i] = self._dot(self._tree[i << 1], self._tree[i << 1 | 1])

    def add(self, i, v):
        self.update(i, self[i] + v)

    def sum(self, l, r):
        l += self._n
        r += self._n
        l_val = r_val = self._unit
        while l < r:
            if l & 1:
                l_val = self._dot(l_val, self._tree[l])
                l += 1
            if r & 1:
                r -= 1
                r_val = self._dot(self._tree[r], r_val)
            l >>= 1
            r >>= 1
        return self._dot(l_val, r_val)

from operator import add
def resolve():
    n = int(input())
    color = [[] for _ in range(n)]
    for v, c in enumerate(map(int, input().split())):
        color[c - 1].append(v)
    E = [[] for _ in range(n)]
    for _ in range(n - 1):
        u, v = map(int, input().split())
        u -= 1; v -= 1
        E[u].append(v)
        E[v].append(u)

    # Euler tour
    start = [-1] * n
    end = [-1] * n
    count = 0
    stack = [(~0, -1), (0, -1)]
    while stack:
        v, p = stack.pop()
        if v >= 0:
            start[v] = count
            count += 1
            for nv in E[v]:
                if nv == p:
                    continue
                stack.append((~nv, v))
                stack.append((nv, v))
        else:
            end[~v] = count

    f = lambda x : x * (x + 1) // 2
    ans = [f(n)] * n
    tree = SegmentTree([1] * n, add, 0)
    history = []
    for c in range(n):
        for v in sorted(color[c], key = lambda x : start[x], reverse = 1):
            count = 1
            for nv in E[v]:
                if start[v] > start[nv]: # nv == p
                    continue
                s = tree.sum(start[nv], end[nv])
                ans[c] -= f(s)
                count += s
            tree.add(start[v], -count)
            history.append((start[v], count))
        ans[c] -= f(tree.sum(0, n))
        while history:
            tree.add(*history.pop())

    print(*ans, sep = '\n')
resolve()

Submission status: Accepted

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

n = int(input())
c = list(map(lambda x: int(x) - 1, input().split()))

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

ans = [n * (n + 1) // 2] * n
n_subtree = [1] * n  # n_st[i] = (頂点iの部分木の頂点数)
cnt = [0] * n  # cnt[i] = (訪問済み頂点のうち、色iを通過しないとたどり着けない頂点数)
n_visited = 0


def dfs(v, v_p):
    global n_visited
    cnt_v_before = cnt[c[v]]
    for v_next in graph[v]:
        if v_next == v_p:
            continue

        m_prev = n_visited - cnt[c[v]]
        dfs(v_next, v)
        n_subtree[v] += n_subtree[v_next]
        m_next = n_visited - cnt[c[v]]

        m = m_next - m_prev
        ans[c[v]] -= m * (m + 1) // 2

    cnt[c[v]] = cnt_v_before + n_subtree[v]
    n_visited += 1


dfs(0, -1)

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

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

Submission status: Accepted

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

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

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

Submission status: Accepted

Solution code:
def main():
    import sys
    sys.setrecursionlimit(10**9)
    input = sys.stdin.readline

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

    # 1回以上 > 全体 - 1回も通らない
    # 選べる数 = n(n+1)//2
    size_subtree = [1] * N
    num_v = [0] * N
    ans = [0] * N

    def dfs(v_now, v_parent, color_parent):
        color_now = C[v_now]
        num_v_now = num_v[color_now]
        num_v_parent = num_v[color_parent]
        for v2 in tree[v_now]:
            if v2 == v_parent: continue
            dfs(v2, v_now, color_now)
            size_subtree[v_now] += size_subtree[v2]
        num_v[color_now] = num_v_now + size_subtree[v_now]
        if color_parent != -1:
            k = size_subtree[v_now] - (num_v[color_parent]-num_v_parent)
            ans[color_parent] += k*(k+1)//2

    dfs(0, -1, -1)

    for color in range(N):
        if color == C[0]: continue
        k = N - num_v[color]
        ans[color] += k*(k+1)//2

    numAll = N*(N+1)//2
    ans = [numAll-ans for ans in ans]

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

main()

Submission status: Accepted

Solution code:
import sys

input=sys.stdin.readline
sys.setrecursionlimit(10**7)

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

subtree=[1]*N
def size(v,pv):
    for nv in edge[v]:
        if nv!=pv:
            size(nv,v)
            subtree[v]+=subtree[nv]

size(0,-1)

data=[[] for i in range(N)]
Parent=[[0] for i in range(N)]
parent=[0]*N
def dfs(v,pv,nop):
    pp=parent[nop]
    parent[nop]=v
    Parent[nop].append(v)
    data[c[v]].append((parent[c[v]],v))
    for nv in edge[v]:
        if nv!=pv:
            dfs(nv,v,c[v])
    parent[nop]=pp

dfs(0,-1,0)

for i in range(N):
    dic={v:subtree[v] for v in Parent[i]}
    for p,ch in data[i]:
        dic[p]-=subtree[ch]

    res=N*(N+1)//2
    for p in dic:
        res-=dic[p]*(dic[p]+1)//2
    print(res)

Submission status: Accepted

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

        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: Accepted

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

        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: Accepted

Solution code:
import sys


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


n = int(input())
c = list(map(int, 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 = tuple(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)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]
# print(ans)


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

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

Submission status: Accepted

Solution code:
from typing import List
import sys

input = lambda: sys.stdin.readline().rstrip()
n = int(input())
c = list(map(int, 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 = tuple(map(int, input().split()))
    a -= 1
    b -= 1
    adj[a].append(b)
    adj[b].append(a)


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


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

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

Submission status: Accepted

Solution code:
from typing import List
import sys

n = int(input())
c = list(map(int, 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 = tuple(map(int, input().split()))
    a -= 1
    b -= 1
    adj[a].append(b)
    adj[b].append(a)


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


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

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

Submission status: Accepted

Solution code:
import sys
n = int(input())
c = list(map(int, 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 = tuple(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)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]
# print(ans)


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

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

Submission status: Accepted

Solution code:
import sys
n = int(input())
c = list(map(int, 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 = tuple(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)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]
# print(ans)


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

    color = c[root]
    s = stack[color]
    for child in adj[root]:
        if parent == child:
            continue

        dfs(root, child)
        size[root] += size[child]

        cnt = size[child]
        while s:
            x = s[-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                s.pop()
            else:
                break
        ans[color] -= 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])

Submission status: Accepted

Solution code:
import sys
n = int(input())
c = list(map(int, 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 = tuple(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)]
timer = 0
total = n*(n+1)//2
ans = [total for color in range(n)]
# print(ans)


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

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

        color = c[root]
        cnt = size[child]
        s = stack[c[root]]
        while s:
            x = s[-1]
            if in_time[x] > in_time[root] and out_time[x] != -1:
                cnt -= size[x]
                s.pop()
            else:
                break
        ans[color] -= 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])

Submission status: Accepted

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

def solve():
    N = int(input())
    Cs = [A-1 for A in map(int, input().split())]
    adjL = [[] for _ in range(N)]
    for _ in range(N-1):
        a, b = map(int, input().split())
        a, b = a-1, b-1
        adjL[a].append(b)
        adjL[b].append(a)

    sizeSubtrees = [1] * N
    numVs = [0] * N
    anss = [0] * N

    def dfs(vNow, vPar, clrPar):
        clrNow = Cs[vNow]
        numVNow = numVs[clrNow]
        numVPar = numVs[clrPar]
        for v2 in adjL[vNow]:
            if v2 == vPar: continue
            dfs(v2, vNow, clrNow)
            sizeSubtrees[vNow] += sizeSubtrees[v2]
        numVs[clrNow] = numVNow + sizeSubtrees[vNow]
        if clrPar != -1:
            k = sizeSubtrees[vNow] - (numVs[clrPar]-numVPar)
            anss[clrPar] += k*(k+1)//2

    dfs(0, -1, -1)

    for clr in range(N):
        if clr == Cs[0]: continue
        k = N - numVs[clr]
        anss[clr] += k*(k+1)//2

    numAll = N*(N+1)//2
    anss = [numAll-ans for ans in anss]

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


solve()

Submission status: Accepted

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]
for x in vs_dfs[::-1]:
    for u in G2[x]:
        sz[x] += sz[u]

stack = [1]
path = []
R = [0] * (N+1)
kh = [[N] 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]
        tmp.append(kh[X[v]][-1]-1-sum(tmp))
        kh[X[v]][-1] -= sz[v]
        R[X[v]] += f(tmp)
    else:
        c = path.pop()
        kh[X[-v]].pop()
for i in range(1, N+1):
    print(R[i])