Problem:


Score : 600 points

Problem Statement

We have a tree with N vertices numbered 1 to N. The i-th edge in this tree connects Vertex a_i and b_i.
Additionally, each vertex is painted in a color, and the color of Vertex i is c_i. Here, the color of each vertex is represented by an integer between 1 and N (inclusive). The same integer corresponds to the same color; different integers correspond to different colors.

For each k=1, 2, ..., N, solve the following problem:

Find the number of simple paths that visit a vertex painted in the color k one or more times.

Note: The simple paths from Vertex u to v and from v to u are not distinguished.

Constraints

1 \leq N \leq 2 \times 10^5

1 \leq c_i \leq N

1 \leq a_i,b_i \leq N

The given graph is a tree.

All values in input are integers.

Input

Input is given from Standard Input in the following format:

N
c_1 c_2 ... c_N
a_1 b_1
:
a_{N-1} b_{N-1}

Output

Print the answers for k = 1, 2, ..., N in order, each in its own line.

Sample Input 1

3
1 2 1
1 2
2 3

Sample Output 1

5
4
0

Let P_{i,j} denote the simple path connecting Vertex i and j.

There are 5 simple paths that visit a vertex painted in the color 1 one or more times:

P_{1,1}\,,\,
P_{1,2}\,,\,
P_{1,3}\,,\,
P_{2,3}\,,\,
P_{3,3} 

There are 4 simple paths that visit a vertex painted in the color 2 one or more times:

P_{1,2}\,,\,
P_{1,3}\,,\,
P_{2,2}\,,\,
P_{2,3} 

There are no simple paths that visit a vertex painted in the color 3 one or more times.  

Sample Input 2

1
1

Sample Output 2

1

Sample Input 3

2
1 2
1 2

Sample Output 3

2
2

Sample Input 4

5
1 2 3 4 5
1 2
2 3
3 4
3 5

Sample Output 4

5
8
10
5
5

Sample Input 5

8
2 7 2 5 4 1 7 5
3 1
1 2
2 7
4 5
5 6
6 8
7 8

Sample Output 5

18
15
0
14
23
0
23
0



Sample input:
3
1 2 1
1 2
2 3


Sample output:
5
4
0

Submission status: Wrong Answer

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

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

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

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

    NORMAL = 0
    REFRESH = 1

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

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

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

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

    def save(self):
        self.history = []

    def refresh(self):
        self.status = self.REFRESH
        for (pos, val) in self.history:
            self.add(pos, -1 * val)
        self.status = self.NORMAL
        self.history = []

class Graph:

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

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

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

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


class Solver:

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

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

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

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

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

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

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

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

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









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

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

Submission status: Wrong Answer

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

sys.setrecursionlimit(210000)


class Node(object):

  def __init__(self, idx, color):
    self.idx = idx
    self.color = color
    self.edges = []
    self.traced = False

  def trace(self, counts, excludes):
    self.traced = True
    num = 1

    for edge in self.edges:
      if edge.traced:
        continue

      excludes[self.color].append(0)
      n = edge.trace(counts, excludes)
      p = n - excludes[self.color].pop()
      counts[self.color] += p * (p + 1) // 2
      num += n

    excludes[self.color][-1] = num

    return num


def main():
  N = int(input())
  nodes = [Node(i, v - 1) for i, v in enumerate(map(int, input().split()))]
  counts = [0] * N
  excludes = [[0] for _ in range(N)]

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

  num = nodes[0].trace(counts, excludes)
  pair = num * (num + 1) // 2

  for i in range(len(counts)):
    p = num - excludes[i][0]
    counts[i] += p * (p + 1) // 2

  for c in counts:
    print(pair - c)


if __name__ == '__main__':
  main()

Submission status: Wrong Answer

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

sys.setrecursionlimit(210000)


class Node(object):

  def __init__(self, idx, color):
    self.idx = idx
    self.color = color
    self.edges = []
    self.traced = False

  def trace(self, counts, excludes):
    self.traced = True
    num = 1

    for edge in self.edges:
      if edge.traced:
        continue

      excludes[self.color].append(0)
      n = edge.trace(counts, excludes)
      p = n - excludes[self.color].pop()
      counts[self.color] += p * (p - 1) // 2 + p
      num += n

    excludes[self.color][-1] = num

    return num


def main():
  N = int(input())
  nodes = [Node(i, v - 1) for i, v in enumerate(map(int, input().split()))]
  counts = [0] * N
  excludes = [[0] for _ in range(N)]

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

  num = nodes[0].trace(counts, excludes)
  pair = num * (num - 1) // 2 + num

  for i in range(len(counts)):
    p = num - excludes[i][0]
    counts[i] += p * (p - 1) // 2 + p

  for c in counts:
    print(pair - c)


if __name__ == '__main__':
  main()

Submission status: Wrong Answer

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


class Node(object):

  def __init__(self, idx, color):
    self.idx = idx
    self.color = color
    self.edges = []
    self.traced = False

  def trace(self, counts, excludes):
    self.traced = True
    num = 1

    for edge in self.edges:
      if edge.traced:
        continue

      #excludes[self.color].append(0)
      #n = edge.trace(counts, excludes)
      #p = n - excludes[self.color].pop()
      #counts[self.color] += p * (p - 1) // 2 + p
      #num += n

    #excludes[self.color][-1] = num

    return num


def main():
  N = int(input())
  nodes = [Node(i, v - 1) for i, v in enumerate(map(int, input().split()))]
  counts = [0] * N
  excludes = [[0] for _ in range(N)]

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

  num = nodes[0].trace(counts, excludes)
  pair = num * (num - 1) // 2 + num

  for i in range(len(counts)):
    p = num - excludes[i][0]
    counts[i] += p * (p - 1) // 2 + p

  for c in counts:
    print(pair - c)


if __name__ == '__main__':
  main()

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
import sys
sys.setrecursionlimit(10**8)
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))


def merge(d: dict, e: dict):
    if len(d) < len(e):
        d, e = e, d
    for k in e:
        d[k] += e[k]
    return e


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)

ans = [0]*N


def f(curr, par=-1):
    ret = defaultdict(int)
    size = 1
    for dest, w in g.E[curr]:
        if dest == par:
            continue
        sz, child = f(dest, curr)
        size += sz

        n = sz-child[c[curr]]
        ans[c[curr]] += n*(n+1)//2

        # マージ
        # if len(ret) < len(child):
        #     child, ret = ret, child
        # for key in child:
        #     ret[key] += child[key]
        ret = merge(ret, child)

    ret[c[curr]] = size
    return size, ret


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

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

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
import sys
sys.setrecursionlimit(10**8)
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))


def merge(d: dict, e: dict):
    if len(d) < len(e):
        d, e = e, d
    for k in d:
        d[k] += e[k]
    return e


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)

ans = [0]*N


def f(curr, par=-1):
    ret = defaultdict(int)
    size = 1
    for dest, w in g.E[curr]:
        if dest == par:
            continue
        sz, child = f(dest, curr)
        size += sz

        n = sz-child[c[curr]]
        ans[c[curr]] += n*(n+1)//2

        # マージ
        # if len(ret) < len(child):
        #     child, ret = ret, child
        # for key in child:
        #     ret[key] += child[key]
        ret = merge(ret, child)

    ret[c[curr]] = size
    return size, ret


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

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

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
import sys
sys.setrecursionlimit(10**8)
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))


def merge(d: dict, e: dict):
    if len(d) > len(e):
        d, e = e, d
    for k in d:
        d[k] += e[k]
    return e


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)

ans = [0]*N


def f(curr, par=-1):
    ret = defaultdict(int)
    size = 1
    for dest, w in g.E[curr]:
        if dest == par:
            continue
        sz, child = f(dest, curr)
        size += sz

        n = sz-child[c[curr]]
        ans[c[curr]] += n*(n+1)//2

        ret = merge(ret, child)

    ret[c[curr]] = size
    return size, ret


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

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

Submission status: Wrong Answer

Solution code:
#!/usr/bin/env python3
import sys
sys.setrecursionlimit(10**8)
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):
        self.E[f].append(t)
        self.E[t].append(f)

    def add_edges(self, A, B, offset=-1):
        for a, b in zip(A, B):
            self.add_edge(a+offset, b+offset)

    def gen_child(self, node, par=-1):
        for child in self.E[node]:
            if child == par:
                continue
            yield child


def make_order(g, v):
    seen = [False]*g.N
    last_order = [-1]*g.N
    parent = [-1]*g.N
    size = [1]*g.N
    counter = {"last": 0}

    def recur(v):
        seen[v] = True
        for to in g.E[v]:
            if seen[to] == True:
                continue
            parent[to] = v
            recur(to)
            size[v] += size[to]

        last_order[counter["last"]] = v
        counter["last"] += 1

    recur(v)
    return last_order, parent, size


def merge(d: dict, e: dict):
    if len(d) > len(e):
        d, e = e, d
    for k in d:
        d[k] += e[k]
    return e


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)
g.add_edges(A, B)


ans = [0]*N
last_order, parent, size = make_order(g, 0)

ret = {}


def f(x, n):
    a = n-x
    return a*(a+1)//2


for curr in last_order:
    ret[curr] = defaultdict(int)
    for dest in g.gen_child(curr, parent[curr]):
        child = ret[dest]
        ans[c[curr]] += f(child[c[curr]], size[dest])
        ret[curr] = merge(ret[curr], child)
    ret[curr][c[curr]] = size[curr]

for color in range(N):
    if color != c[0]:
        ans[color] += f(ret[0][color], N)

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

Submission status: Wrong Answer

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を通らないパスの総数
minus = [0] * (n+1) # uの部分木のうち
minus0 = defaultdict(int)
st = {k: [] for k in range(1, n+1)}
nextu = {}
def dfs(u, prev, ans):
    global minus, st, minus0
    k = cs[u]
    st[k].append(u)
    ss = 1
    for v in ns[u]:
        if v==prev:
            continue
        nextu[u] = v
        s = dfs(v, u, ans)
        ss += s
        tmp = s - minus[v]
        ans[k] += (int(tmp * (tmp - 1) / 2) + tmp)
    st[k].pop()
    if st[k]:
        minus[nextu[st[k][-1]]] += ss
    else:
        minus0[k] += ss
#     print(ss)
    return ss

dfs(1, -1, ans)
for k in range(1, n+1):
    tmp = n - minus0[k]
    ans[k] += (int(tmp * (tmp - 1) / 2) + tmp)
    
total = int(n * (n+1) / 2)
for i in range(1, n+1):
    ans[i] = total - ans[i]
print(*ans, sep="\n")

Submission status: Wrong Answer

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

N=int(input())
C=[0]+list(map(int,input().split()))
E=[[] for i in range(N+1)]

for i in range(N-1):
    a,b=map(int,input().split())
    E[a].append(b)
    E[b].append(a)
        
from collections import deque
QUE = deque([1])
QUE2 = deque()
EULER=[] # これが1からツアーで辿った点
USED=[0]*(N+1)
Parent=[-1]*(N+1)
TOP_SORT=[1]

while QUE:
    x=QUE.pop()
    EULER.append(x)
    
    if USED[x]==1:
        continue
    
    for to in E[x]:

        if USED[to]==0:
            QUE2.append(to)
            Parent[to]=x
            TOP_SORT.append(to)
        else:
            QUE.append(to)
    QUE.extend(QUE2)
    QUE2=deque()

    USED[x]=1

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

for e in EULER:
    CE[C[e]].append(e)

Size=[1]*(N+1)
for i in TOP_SORT[1:][::-1]:
    Size[Parent[i]]+=Size[i]

Size[0]=N+1

CFIRST=[-1]*(N+1)
CLAST=[-1]*(N+1)
P2=[-1]*(N+1)
Child=[[] for i in range(N+1)]
for color in range(1,N+1):
    if CE[color]==[]:
        print(0)
        continue

    EULER2=CE[color]
    for i in range(len(EULER2)):
        CLAST[EULER2[i]]=i

        if CFIRST[EULER2[i]]==-1:
            CFIRST[EULER2[i]]=i

    #print(CE[color])
    #print(CFIRST)
    #print(CLAST)
    #print()

    ANS=0
    Q=[0]
    Child[0]=[]
    

    for i in range(len(EULER2)):
        if CFIRST[EULER2[i]]==i:
            P2[EULER2[i]]=Q[-1]
            Child[Q[-1]].append(EULER2[i])
            Q.append(EULER2[i])
        if CLAST[EULER2[i]]==i:
            Q.pop()

    #print(P2)
    #print(Child)

    ANS=0
    Q=[0]
    while Q:
        x=Q.pop()
        S=Size[x]-1
        for to in Child[x]:
            Q.append(to)
            S-=Size[to]
        #print(Child[x],x,S)
        ANS+=S*(S+1)//2

    print(N*(N+1)//2-ANS)

Submission status: Wrong Answer

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


def EulerTour(graph, root=1):
    N = len(graph) - 1
    parent = [0] * (N + 1)
    stack = [-root, root]
    tour = []
    ind_L = [0] * (N + 1)
    ind_R = [0] * (N + 1)
    i = -1
    while stack:
        v = stack.pop()
        tour.append(v)
        i += 1
        if v > 0:
            ind_L[v] = i
            p = parent[v]
            for w in graph[v]:
                if w == p:
                    continue
                parent[w] = v
                stack.append(-w)
                stack.append(w)
        else:
            ind_R[-v] = i
    return tour, ind_L, ind_R, parent


N = int(readline())
graph = [[] for _ in range(N + 1)]
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
for a, b in zip(m, m):
    graph[a].append(b)
    graph[b].append(a)

tour, ind_L, ind_R, parent = EulerTour(graph)

Submission status: Wrong Answer

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)
    sys.stdout.write(str(item) + "\n")

Submission status: Wrong Answer

Solution code:
mod = 1000000007
eps = 10**-9


def main():
    import sys
    from collections import deque
    from bisect import bisect_left
    input = sys.stdin.buffer.readline

    def f(n):
        return (n * (n-1)) // 2 + n

    # adj[0] must be empty list
    def EulerTour(adj, root):
        st = [root]
        ret = []
        seen = [0] * len(adj)
        par = [0] * len(adj)
        depth = [0] * len(adj)
        while st:
            v = st.pop()
            if seen[v]:
                ret.append(v)
                continue
            ret.append(v)
            seen[v] = 1
            if par[v] != 0:
                st.append(par[v])
            for u in adj[v]:
                if seen[u] == 0:
                    st.append(u)
                    par[u] = v
                    depth[u] = depth[v] + 1

        return ret, depth

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

    que = deque()
    que.append(1)
    seen = [-1] * (N+1)
    seen[1] = 0
    par = [0] * (N+1)
    child = [[] for _ in range(N+1)]
    seq = []
    while que:
        v = que.popleft()
        seq.append(v)
        for u in adj[v]:
            if seen[u] == -1:
                seen[u] = seen[v] + 1
                par[u] = v
                child[v].append(u)
                que.append(u)
    seq.reverse()

    dp = [1] * (N+1)
    for v in seq:
        for u in child[v]:
            dp[v] += dp[u]

    et, _ = EulerTour(adj, 1)
    C_v = [[] for _ in range(N+1)]
    C_idx = [[] for _ in range(N+1)]
    left = [-1] * (N+1)
    for i, v in enumerate(et):
        c = C[v-1]
        C_v[c].append(v)
        C_idx[c].append(i)
        if left[v] < 0:
            left[v] = i
    ans = [f(N)] * (N+1)
    sub = [0] * (N+1)
    sub_0 = [0] * (N+1)
    for v in seq:
        c = C[v-1]
        for u in child[v]:
            ans[c] -= f(dp[u] - sub[u])
            #print(v, u, ans)
        u_idx = bisect_left(C_idx[c], left[v])
        if u_idx:
            w_idx = C_idx[c][u_idx-1] + 1
            w = et[w_idx]
            sub[w] += dp[v]
        else:
            sub_0[c] += dp[v]
    for c in range(1, N+1):
        print(ans[c] - f(N - sub_0[c]))


if __name__ == '__main__':
    main()

Submission status: Wrong Answer

Solution code:
import collections
import sys

sys.setrecursionlimit(10 ** 6)

N = int(input())
colors = [int(x) - 1 for x in input().split()]
neighbors = collections.defaultdict(list)
for i in range(N - 1):
    a, b = map(int, input().split())
    neighbors[a - 1].append(b - 1)
    neighbors[b - 1].append(a - 1)

subtree_sizes = [0] * N


def calc_sizes(node, parent):
    sz = 1
    for i in neighbors[node]:
        if i != parent:
            sz += calc_sizes(i, node)
    subtree_sizes[node] = sz
    return sz


calc_sizes(0, -1)
components = {c: [N] for c in range(N)}
components[colors[0]] = [0]


def calc_components(node, parent):
    color = colors[node]
    if parent >= 0 and colors[parent] != color:
        components[colors[parent]].append(subtree_sizes[node])
        components[color][-1] -= subtree_sizes[node]
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node)


calc_components(0, -1)


def count_paths(n):
    return n * (n + 1) // 2


total_paths = count_paths(subtree_sizes[0])
for c in range(N):
    print(total_paths - sum(count_paths(size) for size in components[c]))

Submission status: Wrong Answer

Solution code:
import collections
import sys

sys.setrecursionlimit(10 ** 6)

N = int(input())
colors = [int(x) - 1 for x in input().split()]
neighbors = collections.defaultdict(list)
for i in range(N - 1):
    a, b = map(int, input().split())
    neighbors[a - 1].append(b - 1)
    neighbors[b - 1].append(a - 1)

subtree_sizes = [0] * N


def calc_sizes(node, parent):
    sz = 1
    for i in neighbors[node]:
        if i != parent:
            sz += calc_sizes(i, node)
    subtree_sizes[node] = sz
    return sz


calc_sizes(0, -1)
components = {c: [N] for c in range(N)}


def calc_components(node, parent):
    color = colors[node]
    components[color][-1] -= subtree_sizes[node]
    if parent >= 0 and colors[parent] != color:
        components[colors[parent]].append(subtree_sizes[node])
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node)


calc_components(0, -1)


def count_paths(n):
    return n * (n + 1) // 2


total_paths = count_paths(subtree_sizes[0])
for c in range(N):
    print(total_paths - sum(count_paths(size) for size in components[c]))

Submission status: Wrong Answer

Solution code:
import collections
import sys

sys.setrecursionlimit(10 ** 6)

N = int(input())
colors = [int(x) - 1 for x in input().split()]
neighbors = collections.defaultdict(list)
for i in range(N - 1):
    a, b = map(int, input().split())
    neighbors[a - 1].append(b - 1)
    neighbors[b - 1].append(a - 1)

subtree_sizes = [0] * N


def calc_sizes(node, parent):
    sz = 1
    for i in neighbors[node]:
        if i != parent:
            sz += calc_sizes(i, node)
    subtree_sizes[node] = sz
    return sz


calc_sizes(0, -1)
components = {c: [subtree_sizes[0]] for c in range(N)}


def calc_components(node, parent):
    color = colors[node]
    if parent >= 0 and colors[parent] != color:
        components[colors[parent]].append(subtree_sizes[node])
    components[color][-1] -= subtree_sizes[node]
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node)


calc_components(0, -1)


def count_paths(n):
    return n * (n + 1) // 2


total_paths = count_paths(subtree_sizes[0])
for c in range(N):
    print(total_paths - sum(count_paths(size) for size in components[c]))

Submission status: Wrong Answer

Solution code:
n = int(input())
c = list(map(int, input().split()))
ab = [list(map(int, input().split())) for _ in range(n - 1)]


def comb2(x):
    if x < 2:
        return 0

    return x * (x - 1) // 2


adj = [[] for _ in range(n)]
for a, b in ab:
    a -= 1
    b -= 1
    adj[a].append(b)
    adj[b].append(a)

s = 0
stack = [s]
d = [-1] * n
d[s] = 0
par = [-1] * n
c = [[] for _ in range(n)]
order = []
while stack:
    u = stack.pop()
    order.append(u)
    for v in adj[u]:
        if d[v] == -1:
            d[v] = d[u] + 1
            par[v] = u
            c[u].append(v)
            stack.append(v)

size = [0] * n
for v in order[::-1]:
    p = par[v]
    if p != -1:
        size[p] += size[v] + 1

for v in range(n):
    ans = 1
    cnt = 1
    for e in c[v]:
        ans *= size[e] + 2
        cnt += 1
    p = par[v]
    if p != -1:
        ans *= n - size[p] + 1
    ans -= comb2(cnt)
    ans += 1

    print(ans)

Submission status: Wrong Answer

Solution code:
from copy import deepcopy
n=int(input())
col=list(map(int,input().split()))
children=[[] for i in range(n)]
n_path=[0]*n
c_path=[[0]*n for i in range(n)]
checked=[False]*n
num_check=0
is_inv=False
for i in range(n):
  col[i]-=1
for i in range(n-1):
  a,b=map(int,input().split())
  children[a-1].append(b-1)
for i in children:
  for j in range(len(i)):
    if(checked[j]):
      is_inv=True
      break
    checked[j]=True
if(is_inv):
  tmp_children=deepcopy(children)
  children=[[] for i in range(n)]
  for i in range(n):
    for j in range(len(tmp_children[i])):
      children[tmp_children[i][j]].append(i)
checked=[False]*n
while num_check<n:
  for parent in range(n):
    if(checked[parent]):
      continue
    computable=True
    for child in children[parent]:
      if(not checked[child]):
        computable=False
        break
    if(not computable):
      continue
    checked[parent]=True
    num_check+=1
    for child in children[parent]:
      n_path[parent]+=n_path[child]
      for col_i in range(n):
        c_path[parent][col_i]+= (n_path[child] if col[parent]==col_i else c_path[child][col_i])
    n_path[parent]+=1
    c_path[parent][col[parent]]+=1
for col_i in range(n):
  ans=0
  for p in range(n):
    if(col[p]==col_i):
      ans+=n_path[p]
      for c1 in children[p]:
        for c2 in children[p]:
          if(c1>=c2):
            continue
          ans+=n_path[c1]*n_path[c2]
    else:
      ans+=c_path[p][col_i]
      for c1 in children[p]:
        for c2 in children[p]:
          if(c1==c2):
            continue
          ans+=c_path[c1][col_i]*(n_path[c2] if c2>c1 else n_path[c2]-c_path[c2][col_i])
  print(ans)

Submission status: Wrong Answer

Solution code:
import sys
from collections import defaultdict as dd
input = sys.stdin.readline
N = int(input())
a = list(map(int, input().split()))
e = dd(list)
for _ in range(N - 1):
  u, v = map(int, input().split())
  e[u].append(v)
  e[v].append(u)

s = [1]
ps = [-1] * (N + 1)
order = []
vis = [0] * (N + 1)
while len(s):
  x = s.pop()
  vis[x] = 1
  order.append(x)
  for y in e[x]:
    if vis[y]: continue
    s.append(y)
    ps[y] = x
size = [1] * (N + 1)
for x in order[: : -1]:
  y = ps[x]
  if y == -1: continue
  size[y] += size[x]


res = [0] * (N + 1)
table = [0] * (N + 1)
for x in range(N):
  table[a[x]] += 1
for x in range(N):
  res[a[x]] += size[x + 1] * (N + 1 - size[x + 1])
for r in res[1: ]: print(r)