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

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

path = [set() for _ in range(n)]

for i in range(n-1):
    a,b = map(int,input().split())
    #a,b = map(int,line[i].split())
    a -= 1
    b -= 1
    path[a].add(b)
    path[b].add(a)
    
ans = [n*(n+1)//2]*n
size = [0]*n
cparent = [[] for _ in range(n)]
reached = [False]*n
root_size = [n]*n

def dfs (p):
    c = C[p] - 1
    cparent[c].append(p)
    s = 1
    for nxt in path[p]:
        #print (nxt)
        size[p] = 0
        if reached[nxt]:
            continue
        reached[nxt] = True
        ret = dfs(nxt)
        s += ret
        size[p] +=ret
        ans[c] -= size[p] * (size[p] + 1)//2
    cparent[c].pop()
    if cparent[c]:
        size[cparent[c][-1]] -=s
    else:
        root_size[c] -=s    
    return s

reached[0] = True
dfs (0)

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

Submission status: Runtime Error

Solution code:
#n = 5
#C = list(map(int,"1 2 3 4 5".split()))
#line = ["1 2", "2 3", "3 4", "3 5"]


n = int(input())
C = list(map(int,input().split()))
 
ci = [set() for _ in range(n)]
for i,c in enumerate(C):
    ci[c-1].add(i)

path = [set() for _ in range(n)]
 
#line = ["1 2", "2 3"]
    
for i in range(n-1):
    a,b = map(int,input().split())
    #a,b = map(int,line[i].split())
    a -= 1
    b -= 1
    path[a].add(b)
    path[b].add(a)
    
ans = [n*(n+1)//2]*n
size = [0]*n
cparent = [[] for _ in range(n)]
reached = [False]*n
root_size = [n]*n

def dfs (p):
    c = C[p] - 1
    cparent[c].append(p)
    s = 1
    for nxt in path[p]:
        #print (nxt)
        size[p] = 0
        if reached[nxt]:
            continue
        reached[nxt] = True
        ret = dfs(nxt)
        s += ret
        size[p] +=ret
        ans[c] -= size[p] * (size[p] + 1)//2
    cparent[c].pop()
    if cparent[c]:
        size[cparent[c][-1]] -=s
    else:
        root_size[c] -=s    
    return s

reached[0] = True
dfs (0)

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

Submission status: Runtime Error

Solution code:
#n = 5
#C = list(map(int,"1 2 3 4 5".split()))
#line = ["1 2", "2 3", "3 4", "3 5"]


n = int(input())
C = list(map(int,input().split()))
 
ci = [set() for _ in range(n)]
for i,c in enumerate(C):
    ci[c-1].add(i)

path = [set() for _ in range(n)]
 
#line = ["1 2", "2 3"]
    
for i in range(n-1):
    #a,b = map(int,input().split())
    a,b = map(int,line[i].split())
    a -= 1
    b -= 1
    path[a].add(b)
    path[b].add(a)
    
ans = [n*(n+1)//2]*n
size = [0]*n
cparent = [[] for _ in range(n)]
reached = [False]*n
root_size = [n]*n

def dfs (p):
    c = C[p] - 1
    cparent[c].append(p)
    s = 1
    for nxt in path[p]:
        #print (nxt)
        size[p] = 0
        if reached[nxt]:
            continue
        reached[nxt] = True
        ret = dfs(nxt)
        s += ret
        size[p] +=ret
        ans[c] -= size[p] * (size[p] + 1)//2
    cparent[c].pop()
    if cparent[c]:
        size[cparent[c][-1]] -=s
    else:
        root_size[c] -=s    
    return s

reached[0] = True
dfs (0)

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

Submission status: Runtime Error

Solution code:
import sys
import itertools

sys.setrecursionlimit(3 * 10**5)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.children = []

  def add_child(self, node):
    self.children.append(node)

  @staticmethod
  def prepare_nodes(n):
    nodes = tuple(map(Node, enumerate(input().split())))
    for _ in range(n - 1):
      a, b = map(lambda x: int(x) - 1, input().split())
      nodes[a].add_child(nodes[b])
      nodes[b].add_child(nodes[a])

    return nodes


def main():
  n = int(input())
  nodes = Node.prepare_nodes(n)

  dp = [0 for _ in range(n)]
  ans_list = [n*(n + 1) // 2 for _ in range(n)]

  def update_ans(s, color):
    ans_list[color] -= (s + dp[color]) * (s + dp[color] + 1) // 2

  def dfs(s, node, parent):
    color = node.c
    pre_s = s + dp[color]
    for child in node.children:
      if child == parent: continue

      dp[color] = -s
      s = dfs(s, child, node)
      update_ans(s, color)

    s += 1
    dp[color] = pre_s - s

    return s

  s = dfs(0, nodes[0], Node((-1, -1)))

  for i in range(n):
    update_ans(s, i)

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


main()

Submission status: Runtime Error

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
# 小さすぎても怒られるので
from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, 2*n))

cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)

ans = [0] * (n+1) # kを通らないパスの総数
def dfs(u, prev, ans, results):
    # result : {k: (kを通ってuで終わるパスの数=kを通らないと到達できない頂点数)}, (u以下の部分木のサイズ)
    k = cs[u]
    if len(ns[u])==1 and prev in ns[u]:
        eend, ss = {k : 1}, 1
    else:
        ss = 1
        eend = defaultdict(int)
        for v in ns[u]:
            if v==prev:
                continue
            end, s = results[v, u] #dfs(v, u, ans)
            ss += s
            tmp = s - end.get(k, 0)
            ans[k] += (int(tmp * (tmp - 1) / 2) + tmp)
            for kk,vv in end.items():
                if kk!=k:
                    eend[kk] += end[kk]
        eend[k] = ss
#     print(eend, ss, ans)
    return eend, ss

q = [(1, -1)] # now, prev
args = [None] * n
i = 0
while q:
    u, prev = q.pop()
    args[i] = (u,prev)
    i += 1
    for v in ns[u]:
        if prev==v:
            continue
        q.append((v, u))
results = {}

for u, prev in args[::-1]:
    results[u, prev] =  dfs(u, prev, ans, results)

for k in range(1, n+1):
    if k==cs[0]:
        continue
    tmp = s - end.get(k, 0)
    ans[k] += (int(tmp * (tmp - 1) / 2) + tmp)
    
total = int(n * (n+1) / 2)
for i in range(1, n+1):
    print(total - ans[i])

Submission status: Runtime Error

Solution code:
import sys

sys.setrecursionlimit(4000)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.relations = []

  def add_relation(self, node):
    self.relations.append(node)


def prepare_nodes(n):
  node_dict = {
    node.id: node
    for node in map(Node, enumerate(input().split()))
  }

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

  return node_dict


def get_connected_node_ids(k, node, prev_id = None):
  other_color_nodes = filter(
    lambda node: node.c != k and node.id != prev_id,
    node.relations
  )

  result = {node.id}
  for next_node in other_color_nodes:
    result |= get_connected_node_ids(k, next_node, node.id)
  return result


def main():
  n = int(input())
  node_dict = prepare_nodes(n)
  comb_num = int(n * (n - 1) / 2) + n

  for k in range(1, n + 1):
    other_color_nodes = tuple(filter(
      lambda node: node.c != k, node_dict.values()
    ))
    num = len(other_color_nodes)
    if num == n:
      print(0)
      continue

    done_ids = set()
    for node in other_color_nodes:
      if node.id in done_ids:
        continue

      node_ids = get_connected_node_ids(k, node)
      node_ids_len = len(node_ids)

      num += node_ids_len * (node_ids_len - 1) / 2
      done_ids |= node_ids

    print(comb_num - int(num))


main()

Submission status: Runtime Error

Solution code:
import sys

sys.setrecursionlimit(4000)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.relations = []

  def add_relation(self, node):
    self.relations.append(node)


def prepare_nodes(n):
  node_dict = {
    node.id: node
    for node in map(Node, enumerate(input().split()))
  }

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

  return node_dict


def cal(k, node, prev_id = None):
  other_color_nodes = list(filter(
    lambda node: node.c != k and node.id != prev_id,
    node.relations
  ))
  num = len(other_color_nodes)

  for next_node in other_color_nodes:
    num += cal(k, next_node, node.id)
  return num


def main():
  n = int(input())
  node_dict = prepare_nodes(n)
  comb_num = int(n * (n - 1) / 2) + n

  for k in range(1, n + 1):
    other_color_nodes = tuple(filter(
      lambda node: node.c != k, node_dict.values()
    ))
    num = len(other_color_nodes)
    if num == n:
      print(0)
      continue

    num *= 2
    for node in other_color_nodes:
      num += cal(k, node)

    print(comb_num - int(num / 2))


main()

Submission status: Runtime Error

Solution code:
import sys
from collections import defaultdict

sys.setrecursionlimit(4000)


class Node:
  def __init__(self, tpl):
    self.id = tpl[0] + 1
    self.c = int(tpl[1])
    self.relations = {}

  def add_relation(self, node):
    self.relations[node.id] = node

  def print_node_with_relations(self, nest = 0, exclude_id = None):
    print("  " * nest, self.get_print_str())
    for node in self.relations.values():
      if node.id == exclude_id: continue
      node.print_node_with_relations(nest + 1, self.id)

  def print_node(self):
    print(self.get_print_str())

  def get_print_str(self):
    return f"{self.id}({self.c})"


def prepare_nodes(n):
  node_dict = {
    node.id: node
    for node in map(Node, enumerate(input().split()))
  }

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

  return node_dict


def cal(k, node, prev_id = None):
  other_color_nodes = list(filter(
    lambda node: node.c != k and node.id != prev_id,
    node.relations.values()
  ))
  num = len(other_color_nodes)
  for next_node in other_color_nodes:
    num += cal(k, next_node, node.id)
  return num


def main():
  n = int(input())
  node_dict = prepare_nodes(n)

  for k in range(1, n + 1):
    other_color_nodes = tuple(filter(
      lambda node: node.c != k, node_dict.values()
    ))
    num = len(other_color_nodes) * 2
    if num == n * 2:
      print(0)
      continue

    for node in other_color_nodes:
      num += cal(k, node)

    print(int(n * (n - 1) / 2) + n - int(num / 2))


main()

Submission status: Runtime Error

Solution code:
#!/usr/bin/env python3
import sys
from collections import defaultdict
INF = float("inf")
MOD = 10**9+7


class Graph(object):
    def __init__(self, N):
        self.N = N
        self.E = defaultdict(list)

    def add_edge(self, f, t, w=1):
        self.E[f].append((t, w))
        self.E[t].append((f, w))


N = int(input())
c = [x-1 for x in map(int, input().split())]
A = [None]*(N-1)
B = [None]*(N-1)
for i in range(N-1):
    A[i], B[i] = map(int, input().split())

g = Graph(N)
for a, b in zip(A, B):
    g.add_edge(a-1, b-1)

# k=1, 2, ..., Nに対して
# 色kが塗られている頂点を一度以上通るような単純パスの数を求める
# 全パスの個数 - 色kが塗られている頂点を一度も通らない単純パスの数
# 全パスの個数はN*(N+1)/2

# グラフを色kのノードで分割して、部分グラフ内での単純パスの総数を求めれば良い
# 各ノードに状態として辞書をもたせる。
# 辞書は色iを通らずに到達可能な頂点の数を持つ。

# 辞書
reachable = [defaultdict(int) for _ in range(N)]

# 回答用
ans = [0]*N


def f(curr, par=-1):
    # 再帰関数
    # curr: 現在の節点
    # par : 親節点の番号
    for dest, w in g.E[curr]:
        if dest == par:
            continue
        f(dest, curr)

        for color in range(N):
            n = reachable[dest][color]
            if color == c[curr]:
                # 自身の色と同じ場合、子の頂点の数から加算
                ans[color] += n*(n+1)//2
                reachable[curr][color] = 0
            else:
                # 自身の色以外については、子の情報から自分の情報を作る
                reachable[curr][color] += n
    for color in range(N):
        if color != c[curr]:
            reachable[curr][color] += 1
    return


f(0)
for color in range(N):
    if color != c[0]:
        n = reachable[0][color]
        ans[color] += n*(n+1)//2

tot = N*(N+1)//2
for a in ans:
    print(tot-a)

Submission status: Runtime Error

Solution code:
import sys
import numpy as np
sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))
def resolve():
    N = ir()
    c = lr()
    ab = [[] for i in range(N)]
    for i in range(N-1):
        a, b = lr()
        ab[a-1].append(b-1)
        ab[b-1].append(a-1)
    tot = N*(N+1)//2
    cnt = np.zeros((N, N), np.int32)
    vis = [0]*N
    row = np.arange(N).astype(np.int32)
    col = np.zeros(N, np.int8)
    s = [[0, col]]
    vis[0] += 1
    while len(s) > 0:
        v, col = s.pop()
        cnt[row, col] += 1
        cnt[c[v]-1, col[c[v]-1]] -= 1
        for i in ab[v]:
            if vis[i] == 0:
                col[c[v]-1] += 1
                s.append([i, col.copy()])
                vis[i] += 1
    print(*(tot-np.sum(cnt*(cnt+1)//2, axis=1)).tolist(), sep='\n')
resolve()

Submission status: Runtime Error

Solution code:
import sys
sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))
def resolve():
    N = ir()
    c = lr()
    ab = [[] for i in range(N)]
    for i in range(N-1):
        a, b = lr()
        ab[a-1].append(b-1)
        ab[b-1].append(a-1)
    tot = N*(N+1)//2
    il = [0]*N
    ol = [0]*N
    cnt = [-1]
    def dfs(u, p):
        cnt[0] += 1
        il[u] = cnt[0]
        for i in ab[u]:
            if i != p:
                dfs(i, u)
        ol[u] = cnt[0]
    dfs(0, -1)
    vs = [[] for i in range(N)]
    for i in range(N):
        vs[c[i]-1].append(i)
    for i in range(N):
        ans = tot
        sl = [1]*N
        tvs = sorted(vs[i], key=lambda x: il[x])[::-1]
        for v in tvs:
            n = 1
            for j in ab[v]:
                if il[j] > il[v]:
                    tmp = sum(sl[il[j]:ol[j]+1])
                    ans -= tmp*(tmp+1)//2
                    n += tmp
            sl[il[v]] -= n
        tmp = sum(sl[0:ol[0]+1])
        ans -= tmp*(tmp+1)//2
        print(ans)
resolve()

Submission status: Runtime Error

Solution code:
import sys
import numpy as np
sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))
def resolve():
    N = ir()
    c = lr()
    ab = [[] for i in range(N)]
    for i in range(N-1):
        a, b = lr()
        ab[a-1].append(b-1)
        ab[b-1].append(a-1)
    tot = N*(N+1)//2
    cnt = np.zeros((N, N), np.int8)
    vis = [0]*N
    row = np.arange(N).astype(np.int8)
    col = np.zeros(N, np.int8)
    s = [[0, col]]
    vis[0] += 1
    while len(s) > 0:
        v, col = s.pop()
        cnt[row, col] += 1
        cnt[c[v]-1, col[c[v]-1]] -= 1
        for i in ab[v]:
            if vis[i] == 0:
                col[c[v]-1] += 1
                s.append([i, col.copy()])
                vis[i] += 1
    print(*(tot-np.sum(cnt*(cnt+1)//2, axis=1)).tolist(), sep='\n')
resolve()

Submission status: Runtime Error

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

n = int(input())
c = [int(ci) - 1 for ci in input().split()]

def iig(vn, en):
    res = [[] for _ in range(vn)]
    for _ in range(en):
        ai, bi = map(lambda x: int(x) - 1, input().split())
        res[ai].append(bi)
        res[bi].append(ai)
    return res

def sep_pc(graph, root=0):
    parent = [-1 for _ in range(n)]
    children = [None for _ in range(n)]
    def f(x, p):
        parent[x] = p
        children[x] = [y for y in graph[x] if y != p]
        for y in graph[x]:
            if y != p:
                f(y, x)
        return
    f(root, -1)
    return parent, children

def itree(n, root=0):
    return sep_pc(iig(n, n - 1), root)

parent, children = itree(n)

def get_tsize(children, root=0):
    res = [0 for _ in children]
    def f(x):
        res[x] = sum((f(y) for y in children[x]), 1)
        return res[x]
    f(root)
    return res

tsize = get_tsize(children)
head = [n for _ in range(n)]
location = [-1 for _ in range(n)]
def fill_seg_size(x):
    color = c[x]
    if location[color] < 0:
        head[color] -= tsize[x]
    else:
        tsize[location[color]] -= tsize[x]
    prev_location = location[color]
    for y in children[x]:
        location[color] = y
        fill_seg_size(y)
    location[color] = prev_location
    return

fill_seg_size(0)

def comb2(a):
    return a * (a + 1) // 2

vs_by_color = [[] for _ in range(n)]
for i in range(n):
    vs_by_color[c[i]].append(i)
tot = comb2(n)


for color in range(n):
    xs = vs_by_color[color]
    ans = tot - comb2(head[color]) - sum(comb2(tsize[y]) for x in vs_by_color[color] for y in children[x])
    print(ans)

Submission status: Runtime Error

Solution code:
import sys
sys.setrecursionlimit(10**6+10)

n = int(input())
c = [int(ci) - 1 for ci in input().split()]
v = [[] for _ in range(n)]
for _ in range(n - 1):
    ai, bi = map(lambda x: int(x) - 1, input().split())
    v[ai].append(bi)
    v[bi].append(ai)

par = [-1] * n
ch = [[] for _ in range(n)]
def sep_par_ch(x=0, parent=-1):
    par[x] = parent
    for y in v[x]:
        if y == parent:
            continue
        ch[x].append(y)
        sep_par_ch(y, x)
    return

sep_par_ch()

tsize = [0 for _ in range(n)]
def fill_tsize(x=0):
    tsize[x] = sum(fill_tsize(y) for y in ch[x]) + 1
    return tsize[x]

fill_tsize()

seg_size = [[tsize[y] for y in ch[x]] for x in range(n)]
def fill_seg_size(x, color_ancestor):
    color = c[x]
    alist, aindex = color_ancestor[color]
    alist[aindex] -= tsize[x]
    for i, y in enumerate(ch[x]):
        color_ancestor[color] = seg_size[x], i
        fill_seg_size(y, color_ancestor)
    color_ancestor[color] = alist, aindex
    return

init_ancestor = [([n], 0) for x in range(n)]
fill_seg_size(0, init_ancestor)

def comb2(a):
    return a * (a - 1) // 2 + a

vs_by_color = [[] for _ in range(n)]
for i in range(n):
    vs_by_color[c[i]].append(i)

tot = comb2(n)
for color in range(n):
    xs = vs_by_color[color]
    ans = sum(comb2(w) for x in xs for w in seg_size[x])
    ans += comb2(init_ancestor[color][0][0])
    print(tot - ans)

Submission status: Runtime Error

Solution code:
n = int(input())
c = [int(ci) - 1 for ci in input().split()]
v = [[] for _ in range(n)]
for _ in range(n - 1):
    ai, bi = map(lambda x: int(x) - 1, input().split())
    v[ai].append(bi)
    v[bi].append(ai)

par = [-1] * n
ch = [[] for _ in range(n)]
def sep_parnch(x=0, parent=-1):
    par[x] = parent
    for i in v[x]:
        if i == parent:
            continue
        ch[x].append(i)
        sep_parnch(i, x)
    return

sep_parnch()

tsize = [0] * n
def fill_tsize(root=0):
    tsize[root] = sum(fill_tsize(x) for x in ch[root]) + 1
    return tsize[root]

fill_tsize()

seg_size = [[0] * len(ch[x]) for x in range(n)]
def fill_seg_size(x=0, clist=[(None, -1)]*n):
    al, aind = clist[c[x]]
    if al is not None:
        al[aind] += tsize[x]
    for i, y in enumerate(ch[x]):
        clist[c[x]] = seg_size[x], i
        fill_seg_size(y, clist)
    clist[c[x]] = al, aind
    return

head_seg = [([0], 0) for _ in range(n)]
fill_seg_size(0, head_seg)

def comb2(a):
    return a * (a - 1) // 2 + a

vcat = [[] for _ in range(n)]
for i in range(n):
    vcat[c[i]].append(i)

tot = comb2(n)
for col, vs in enumerate(vcat):
    ans = sum(comb2(tsize[y] - seg_size[x][i]) for x in vs for i, y in enumerate(ch[x])) + comb2(n - head_seg[col][0][0])
    print(tot - ans)

Submission status: Runtime Error

Solution code:
#include<iostream>
#include<vector>
using namespace std;
int N;
int c[2<<17];
vector<int>G[2<<17];
int cnt[2<<17];
long ans[2<<17];
int dfs(int u,int p)
{
	int ret=1;
	int sto=cnt[c[u]];
	int pre=sto;
	for(int v:G[u])if(v!=p)
	{
		int t=dfs(v,u);
		ret+=t;
		long k=t-(cnt[c[u]]-sto);
		ans[c[u]]-=k*(k+1)/2;
		sto=cnt[c[u]];
	}
	cnt[c[u]]=pre+ret;
	return ret;
}
main()
{
	cin>>N;
	for(int i=0;i<N;i++)
	{
		cin>>c[i];
		c[i]--;
	}
	for(int i=1;i<N;i++)
	{
		int u,v;cin>>u>>v;u--,v--;
		G[u].push_back(v);
		G[v].push_back(u);
	}
	dfs(0,-1);
	for(int i=0;i<N;i++)
	{
		ans[i]+=(long)N*(N+1)/2;
		if(i!=c[0])
		{
			long k=N-cnt[i];
			ans[i]-=k*(k+1)/2;
		}
		cout<<ans[i]<<endl;
	}
}

Submission status: Runtime Error

Solution code:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np


def solve(N, C, AB):
    def edges_to_array(A, B, N, M):
        head = np.full(N + 1, -1, np.int32)
        nxt = np.empty(M + M, np.int32)
        to = np.empty(M + M, np.int32)
        for i in range(M):
            a = A[i]
            b = B[i]
            nxt[i << 1] = head[a]
            to[i << 1] = b
            head[a] = i << 1
            nxt[i << 1 | 1] = head[b]
            to[i << 1 | 1] = a
            head[b] = i << 1 | 1
        return head, nxt, to

    def EulerTour(head, nxt, to, root=1):
        N = len(head) - 1
        parent = np.zeros_like(head)
        ind_L = np.empty_like(head)
        ind_R = np.empty_like(head)
        tour = np.empty(N + N, np.int32)
        stack = np.empty(N + N, np.int32)
        stack[0] = -root
        stack[1] = root
        t = -1
        s = 2
        while s > 0:
            s -= 1
            v = stack[s]
            t += 1
            tour[t] = v
            if v > 0:
                ind_L[v] = t
                p = parent[v]
                k = head[v]
                while k != -1:
                    w = to[k]
                    if w == p:
                        k = nxt[k]
                        continue
                    parent[w] = v
                    stack[s] = -w
                    s += 1
                    stack[s] = w
                    s += 1
                    k = nxt[k]
            else:
                ind_R[-v] = t
        return tour, ind_L, ind_R, parent

    head, nxt, to = edges_to_array(AB[::2], AB[1::2], N, N - 1)
    tour, ind_L, ind_R, parent = EulerTour(head, nxt, to)
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


def cc_export():
    from numba.pycc import CC
    cc = CC('my_module')
    cc.export('solve', '(i8,i4[:],i4[:])')
    cc.compile()


if sys.argv[-1] == 'ONLINE_JUDGE':
    cc_export()
    exit()

from my_module import solve
N = int(readline())
C = np.zeros(N + 1, np.int32)
C[1:] = np.array(readline().split(), np.int32)
AB = np.array(read().split(), np.int32)
answer = solve(N, C, AB)
print('\n'.join(answer[1:].astype(str)))

Submission status: Runtime Error

Solution code:
def main():
    n=int(input())
    color=list(map(int,input().split()))
    nodes=[[]for i in range(n)]
    cut=[[0]for i in range(n+1)]
    ans=[0]*(n+1)
    for i in range(n-1):
        a,b=map(int,input().split())
        nodes[a-1].append(b-1)
        nodes[b-1].append(a-1)
    def num(p,parent):
        s=0
        for child in nodes[p]:
            if child==parent:
                continue
            cut[color[p]].append(0)
            nc=num(child,p)
            group=nc-cut[color[p]].pop()
            ans[color[p]]+=group*(group+1)//2
            s+=nc
        s+=1
        cut[color[p]][-1]+=s
        return s
    num(0,-1)
    for a,c in zip(ans[1:],cut[1:]):
        group=n-c[0]
        c=group*(group+1)//2
        print(n*(n+1)//2-c-a)
main()

Submission status: Runtime Error

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, n))
# 小さすぎても怒られるので
 
cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)
 
toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)
 
for item in toru[1:]:
    print(item)

Submission status: Runtime Error

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip() 
from collections import defaultdict

n = int(input())
sys.setrecursionlimit(max(1000, 2*n))
# 小さすぎても怒られるので
 
cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)
 
toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)
 
for item in toru[1:]:
    print(item)
import sys

from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, n))
# 小さすぎても怒られるので

cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)

toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)

for item in toru[1:]:
    print(item)

Submission status: Runtime Error

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip() 
from collections import defaultdict

n = int(input())
sys.setrecursionlimit(max(1000, n))
# 小さすぎても怒られるので
 
cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)
 
toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)
 
for item in toru[1:]:
    print(item)
import sys

from collections import defaultdict
n = int(input())
sys.setrecursionlimit(max(1000, n))
# 小さすぎても怒られるので

cs = [0] + [int(c) for c in input().split()]
ns = defaultdict(set)
for _ in range(n-1):
    a,b = [int(c) for c in input().split()]
    ns[a].add(b)
    ns[b].add(a)

toru = [0] * (n+1)
def dfs(ns, prev, now, toru):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = dfs(ns, now, u, toru)
        # toruの更新
        for kk, vv in tmp_end.items():
            if kk not in end:
                toru[kk] += vv * s
            else:
                toru[kk] += vv * s + end[kk] * (tmp_s+1) - ((vv+1) * end[kk])
        for kk, vv in end.items():
            if kk not in tmp_end:
                toru[kk] += vv * tmp_s
        # end, sの更新
        end[k] += tmp_s
        for kk,vv in tmp_end.items():
            if k!=kk:
                end[kk] += tmp_end[kk]
        s += tmp_s
    toru[k] += 1
#     print(now, prev, end, toru, s)
    return end, toru, s
end, toru, s = dfs(ns, prev=-1, now=1, toru=toru)

for item in toru[1:]:
    print(item)