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:
import sys
sys.setrecursionlimit(2*10**5)

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

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

# 区間データを保管するテーブル
class BIT():
    def __init__(self, n):
        self.size = n
        self.data = [0] * (n+1)  # 今回は1で初期化する。 BIT の先頭は 1からなのでずらす
        
    def add(self, i, s):  # i の情報を更新する
        while i <= self.size:
            self.data[i] += s
            i += i & -i
    
    def sum(self, i):  # 1～i 番目までの合計
        tmp = 0
        while i > 0:
            tmp += self.data[i]
            i -= i & -i
        return tmp

def dfs(u, p = -1):

    global k, I, O
    
    us = 1
    I[u] = k
    k += 1
    
    for v in E[u]:
        if not v == p:
            us += dfs(v, u)
    O[u] = k

    return us

# 経路の組み合わせ数を計算する
def comb2(n):
    return n*(n+1)//2

k = 1
I = [0] * N  # 走査に入ったタイミング
O = [0] * N  # 上に戻ったタイミング
dfs(0)

# 色番号ごとに、頂点をまとめる。まとめらたら、子要素から操作するように並べ替える
VC = [ [] for _ in range(N)]
for j in range(N):
    VC[C[j]].append(j)

for j in range(N):
    VC[j].sort(key=lambda x:I[x], reverse=True)


total = comb2(N)

bit = BIT(N)
for i in range(N):
    bit.add(i+1, 1)

for j in range(N):
        
    ans = total
    history = []
    for u in VC[j]:
        cnt = 1
        for v in E[u]:
            if I[u] > I[v]:
                continue
            num = bit.sum(O[v]-1) - bit.sum(I[v]-1)
            ans -= comb2(num)
            cnt += num
        bit.add(I[u], -cnt)
        history.append((I[u], cnt))
    ans -= comb2(bit.sum(N))
    for a, b in history:
        bit.add(a,b)
    print(ans)

Submission status: Runtime Error

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

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

# 区間データを保管するテーブル
class BIT():
    def __init__(self, n):
        self.size = n
        self.data = [0] * (n+1)  # 今回は1で初期化する。 BIT の先頭は 1からなのでずらす
        
    def add(self, i, s):  # i の情報を更新する
        while i <= self.size:
            self.data[i] += s
            i += i & -i
    
    def sum(self, i):  # 1～i 番目までの合計
        tmp = 0
        while i > 0:
            tmp += self.data[i]
            i -= i & -i
        return tmp

def dfs(u, p = -1):

    global k, I, O
    
    us = 1
    I[u] = k
    k += 1
    
    for v in E[u]:
        if not v == p:
            us += dfs(v, u)
    O[u] = k

    return us

# 経路の組み合わせ数を計算する
def comb2(n):
    return n*(n+1)//2

k = 1
I = [0] * N  # 走査に入ったタイミング
O = [0] * N  # 上に戻ったタイミング
dfs(0)

# 色番号ごとに、頂点をまとめる。まとめらたら、子要素から操作するように並べ替える
VC = [ [] for _ in range(N)]
for j in range(N):
    VC[C[j]].append(j)

for j in range(N):
    VC[j].sort(key=lambda x:I[x], reverse=True)


total = comb2(N)

bit = BIT(N)
for i in range(N):
    bit.add(i+1, 1)

for j in range(N):
        
    ans = total
    history = []
    for u in VC[j]:
        cnt = 1
        for v in E[u]:
            if I[u] > I[v]:
                continue
            num = bit.sum(O[v]-1) - bit.sum(I[v]-1)
            ans -= comb2(num)
            cnt += num
        bit.add(I[u], -cnt)
        history.append((I[u], cnt))
    ans -= comb2(bit.sum(N))
    for a, b in history:
        bit.add(a,b)
    print(ans)

Submission status: Runtime Error

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

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

# 区間データを保管するテーブル
class BIT():
    def __init__(self, n):
        self.size = n
        self.data = [0] * (n+1)  # 今回は1で初期化する。 BIT の先頭は 1からなのでずらす
        
    def add(self, i, s):  # i の情報を更新する
        while i <= self.size:
            self.data[i] += s
            i += i & -i
    
    def sum(self, i):  # 1～i 番目までの合計
        tmp = 0
        while i > 0:
            tmp += self.data[i]
            i -= i & -i
        return tmp

def dfs(u, p = -1):

    global k, I, O
    
    us = 1
    I[u] = k
    k += 1
    
    for v in E[u]:
        if not v == p:
            us += dfs(v, u)
    O[u] = k

    return us

# 経路の組み合わせ数を計算する
def comb2(n):
    return n*(n+1)//2

k = 1
I = [0] * N  # 走査に入ったタイミング
O = [0] * N  # 上に戻ったタイミング
dfs(0)

# 色番号ごとに、頂点をまとめる。まとめらたら、子要素から操作するように並べ替える
VC = [ [] for _ in range(N)]
for j in range(N):
    VC[C[j]].append(j)

for j in range(N):
    VC[j].sort(key=lambda x:I[x], reverse=True)


total = comb2(N)

for j in range(N):
    
for j in range(N):
        
    ans = total
    history = []
    for u in VC[j]:
        cnt = 1
        for v in E[u]:
            if I[u] > I[v]:
                continue
            num = bit.sum(O[v]-1) - bit.sum(I[v]-1)
            ans -= comb2(num)
            cnt += num
        bit.add(I[u], -cnt)
        history.append((I[u], cnt))
    ans -= comb2(bit.sum(N))
    for a, b in history:
        bit.add(a,b)
    print(ans)

Submission status: Runtime Error

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

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

# 区間データを保管するテーブル
class BIT():
    def __init__(self, n):
        self.size = n
        self.data = [0] * (n+1)  # 今回は1で初期化する。 BIT の先頭は 1からなのでずらす
        
    def add(self, i, s):  # i の情報を更新する
        while i <= self.size:
            self.data[i] += s
            i += i & -i
    
    def sum(self, i):  # 1～i 番目までの合計
        tmp = 0
        while i > 0:
            tmp += self.data[i]
            i -= i & -i
        return tmp

def dfs(u, p = -1):

    global k, I, O
    
    us = 1
    I[u] = k
    k += 1
    
    for v in E[u]:
        if not v == p:
            us += dfs(v, u)
    O[u] = k

    return us

# 経路の組み合わせ数を計算する
def comb2(n):
    return n*(n+1)//2

k = 1
I = [0] * N  # 走査に入ったタイミング
O = [0] * N  # 上に戻ったタイミング
dfs(0)

# 色番号ごとに、頂点をまとめる。まとめらたら、子要素から操作するように並べ替える
VC = [ [] for _ in range(N)]
for j in range(N):
    VC[C[j]].append(j)

for j in range(N):
    VC[j].sort(key=lambda x:I[x], reverse=True)


total = comb2(N)

for j in range(N):
    
    bit = BIT(N)
    for i in range(N):
        bit.add(i+1, 1)
    
    ans = total
    for u in VC[j]:
        cnt = 1
        for v in E[u]:
            if I[u] > I[v]:
                continue
            num = bit.sum(O[v]-1) - bit.sum(I[v]-1)
            ans -= comb2(num)
            cnt += num
        bit.add(I[u], -cnt)
    ans -= comb2(bit.sum(N))
    print(ans)

Submission status: Runtime Error

Solution code:
import sys
import socket
from collections import defaultdict

limit = 2 *10 ** 5
sys.setrecursionlimit(limit)

hostnames = ['N551J', 'F551C', 'X553M']
input_file = 'f1.in'
input_file = 'f_random_1_01.txt'
if socket.gethostname() in hostnames:
    sys.stdin = open(input_file)


def read_int_list():
    return list(map(int, input().split()))


def read_int():
    return int(input())


def read_str_list():
    return input().split()


def read_str():
    return input()


def solve_0():
    """
    First attempt.
    O(n^2) solution.
    """
    global c
    n = read_int()
    c = read_int_list()
    a, b = [], []
    if n > 1:
        a, b = map(list, zip(*[read_int_list() for i in range(n - 1)]))
    a = [ x - 1 for x in a]
    b = [ x - 1 for x in b]
    c = [ x - 1 for x in c]

    neighbors = defaultdict(list)
    for i in range(n - 1):
        neighbors[a[i]].append(b[i])
        neighbors[b[i]].append(a[i])

    parent = {}
    children = defaultdict(list)
    n_descendents = defaultdict(int)
    def dfs(i, p=-1):
        parent[i] = p
        n_descendents[i] = 0
        for neighbor in neighbors[i]:
            if neighbor == p:
                continue
            children[i].append(neighbor)
            dfs(neighbor, p=i)
            n_descendents[i] += 1 + n_descendents[neighbor]


    root = 0
    dfs(root, -1)

    # print('parent:', parent)
    # print('children:', children)
    # print('n_descendents:', n_descendents)

    results = []
    for k in range(n):

        # print()
        # print('k:', k)

        n0 = defaultdict(int)
        n1 = defaultdict(int)
        n2 = defaultdict(int)
        def count(i):


            # print('count, i:', i)

            global c
            for child in children[i]:
                count(child)
            n0[i] = int(c[i] == k)

            # n1
            if c[i] == k:
                n1[i] = n_descendents[i]
            else:
                for child in children[i]:
                    n1[i] += n0[child]
                    n1[i] += n1[child]

            # n2
            if c[i] == k:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += (1 + n_descendents[j]) * (1 + n_descendents[l])
            else:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += (1 + n_descendents[j]) * (1 + n_descendents[l])
                        n2[i] -= (1 + n_descendents[j] - n1[j] - n0[j]) * (1 + n_descendents[l] - n1[l] - n0[l])

        count(root)

        # print('n0:', n0)
        # print('n1:', n1)
        # print('n2:', n2)

        res = 0
        for i in range(n):
            res += n0[i] + n1[i] + n2[i]
        results.append(res)

        # print('res:', res)

    return results



def solve():
    n = read_int()
    c = read_int_list()
    a, b = [], []
    if n > 1:
        a, b = map(list, zip(*[read_int_list() for i in range(n - 1)]))
    a = [x - 1 for x in a]
    b = [x - 1 for x in b]
    c = [x - 1 for x in c]

    neighbors = defaultdict(list)
    for i in range(n - 1):
        neighbors[a[i]].append(b[i])
        neighbors[b[i]].append(a[i])

    res = [0] * n
    x = [0] * n
    n_descendents = {}
    def dfs(i, p=-1):

        # print('i:', i)

        n_descendents[i] = 0
        next_x = x[c[i]]
        for neighbor in neighbors[i]:
            if neighbor == p:
                continue
            before = x[c[i]]
            dfs(neighbor, p=i)
            after = x[c[i]]
            s = 1 + n_descendents[neighbor]
            m = s - after + before
            n_descendents[i] += s
            res[c[i]] -= m * (m - 1) // 2

            # if i == 0:

            #     print('neighbor:', neighbor)
            #     print('before:', before)
            #     print('after:', after)
            #     print('s:', s)
            #     print('m:', m)
            #     print('x:', x)
            #     print('res:', res)

        next_x += 1 + n_descendents[i]
        x[c[i]] = next_x
        s = 1 + n_descendents[i]

    # res[c[i]] += s * (s - 1) // 2

    root = 0
    dfs(root, -1)

    # print('n_descendents:', n_descendents)

    # print('x:', x)
    # print('res:', res)

    present_colors = set([x for x in c])

    for i in present_colors:
        if i == c[root]: continue
        m = n - x[i]
        res[i] -= m * (m - 1) // 2
    # print('res:', res)

    # print('n * (n - 1) // 2 : ', n * (n - 1) // 2)


    for i in present_colors:
        res[i] += n * (n - 1) // 2
    # print('res:', res)

    for i in c:
        res[i] += 1
    # print('res:', res)


    return res


def main():
    res = solve()
    print(*res, sep='\n')


if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys
import socket
from collections import defaultdict

hostnames = ['N551J', 'F551C', 'X553M']
input_file = 'f1.in'
# input_file = 'f_random_1_01.txt'
if socket.gethostname() in hostnames:
    sys.stdin = open(input_file)


def read_int_list():
    return list(map(int, input().split()))


def read_int():
    return int(input())


def read_str_list():
    return input().split()


def read_str():
    return input()


def solve_0():
    """
    First attempt.
    O(n^2) solution.
    """
    global c
    n = read_int()
    c = read_int_list()
    a, b = [], []
    if n > 1:
        a, b = map(list, zip(*[read_int_list() for i in range(n - 1)]))
    a = [ x - 1 for x in a]
    b = [ x - 1 for x in b]
    c = [ x - 1 for x in c]

    neighbors = defaultdict(list)
    for i in range(n - 1):
        neighbors[a[i]].append(b[i])
        neighbors[b[i]].append(a[i])

    parent = {}
    children = defaultdict(list)
    n_descendents = defaultdict(int)
    def dfs(i, p=-1):
        parent[i] = p
        n_descendents[i] = 0
        for neighbor in neighbors[i]:
            if neighbor == p:
                continue
            children[i].append(neighbor)
            dfs(neighbor, p=i)
            n_descendents[i] += 1 + n_descendents[neighbor]


    root = 0
    dfs(root, -1)

    # print('parent:', parent)
    # print('children:', children)
    # print('n_descendents:', n_descendents)

    results = []
    for k in range(n):

        # print()
        # print('k:', k)

        n0 = defaultdict(int)
        n1 = defaultdict(int)
        n2 = defaultdict(int)
        def count(i):


            # print('count, i:', i)

            global c
            for child in children[i]:
                count(child)
            n0[i] = int(c[i] == k)

            # n1
            if c[i] == k:
                n1[i] = n_descendents[i]
            else:
                for child in children[i]:
                    n1[i] += n0[child]
                    n1[i] += n1[child]

            # n2
            if c[i] == k:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += (1 + n_descendents[j]) * (1 + n_descendents[l])
            else:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += (1 + n_descendents[j]) * (1 + n_descendents[l])
                        n2[i] -= (1 + n_descendents[j] - n1[j] - n0[j]) * (1 + n_descendents[l] - n1[l] - n0[l])

        count(root)

        # print('n0:', n0)
        # print('n1:', n1)
        # print('n2:', n2)

        res = 0
        for i in range(n):
            res += n0[i] + n1[i] + n2[i]
        results.append(res)

        # print('res:', res)

    return results



def solve():
    n = read_int()
    c = read_int_list()
    a, b = [], []
    if n > 1:
        a, b = map(list, zip(*[read_int_list() for i in range(n - 1)]))
    a = [x - 1 for x in a]
    b = [x - 1 for x in b]
    c = [x - 1 for x in c]

    neighbors = defaultdict(list)
    for i in range(n - 1):
        neighbors[a[i]].append(b[i])
        neighbors[b[i]].append(a[i])

    res = [0] * n
    x = [0] * n
    n_descendents = {}
    def dfs(i, p=-1):

        # print('i:', i)

        n_descendents[i] = 0
        next_x = x[c[i]]
        for neighbor in neighbors[i]:
            if neighbor == p:
                continue
            before = x[c[i]]
            dfs(neighbor, p=i)
            after = x[c[i]]
            s = 1 + n_descendents[neighbor]
            m = s - after + before
            n_descendents[i] += s
            res[c[i]] -= m * (m - 1) // 2

            # if i == 0:

            #     print('neighbor:', neighbor)
            #     print('before:', before)
            #     print('after:', after)
            #     print('s:', s)
            #     print('m:', m)
            #     print('x:', x)
            #     print('res:', res)

        next_x += 1 + n_descendents[i]
        x[c[i]] = next_x
        s = 1 + n_descendents[i]

    # res[c[i]] += s * (s - 1) // 2

    root = 0
    dfs(root, -1)

    # print('n_descendents:', n_descendents)

    # print('x:', x)
    # print('res:', res)

    present_colors = set([x for x in c])

    for i in present_colors:
        if i == c[root]: continue
        m = n - x[i]
        res[i] -= m * (m - 1) // 2
    # print('res:', res)

    # print('n * (n - 1) // 2 : ', n * (n - 1) // 2)


    for i in present_colors:
        res[i] += n * (n - 1) // 2
    # print('res:', res)

    for i in c:
        res[i] += 1
    # print('res:', res)


    return res


def main():
    res = solve()
    print(*res, sep='\n')


if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys
input = sys.stdin.readline
N = int(input())
C = list(map(int,input().split()))
AB = [tuple(map(int,input().split())) for i in range(N-1)]

es = [[] for _ in range(N)]
for a,b in AB:
    a,b = a-1,b-1
    es[a].append(b)
    es[b].append(a)
cs = [[] for _ in range(N)]
for i,c in enumerate(C):
    cs[c-1].append(i)

tin = [-1]*N
tout = [-1]*N
k = 0
def dfs(v,p=-1):
    global k
    tin[v] = k
    k += 1
    for to in es[v]:
        if to==p: continue
        dfs(to,v)
    tout[v] = k
dfs(0)

class BinaryIndexedTree:
    def __init__(self,size):
        self.N = size
        self.bit = [0]*(size+1)
    def add(self,x,w): # 0-indexed
        x += 1
        while x <= self.N:
            self.bit[x] += w
            x += (x & -x)
    def _sum(self,x): # 1-indexed
        ret = 0
        while x > 0:
            ret += self.bit[x]
            x -= (x & -x)
        return ret
    def sum(self,l,r): # [l,r)
        return self._sum(r) - self._sum(l)
    def __str__(self): # for debug
        arr = [self.sum(i,i+1) for i in range(self.N)]
        return str(arr)
bit = BinaryIndexedTree(N)
for i in range(N):
    bit.add(i,1)

whole = N*(N+1)//2
anss = []
for i in range(N):
    ans = whole
    cs[i].sort(key=lambda x: -tin[x])
    history = []
    for v in cs[i]:
        cnt = 1
        for to in es[v]:
            if tin[to] < tin[v]: continue
            num = bit.sum(tin[to], tout[to])
            ans -= num*(num+1)//2
            cnt += num
        bit.add(tin[v], -cnt)
        history.append((tin[v], cnt))
    pn = bit.sum(0,N)
    ans -= pn*(pn+1)//2
    for a,b in history:
        bit.add(a,b)
    anss.append(ans)
print(*anss, sep='\n')

Submission status: Runtime Error

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

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

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

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

Solution code:
import sys
import socket
from collections import defaultdict

hostnames = ['N551J', 'F551C', 'X553M']
input_file = 'f1.in'
if socket.gethostname() in hostnames:
    sys.stdin = open(input_file)


def read_int_list():
    return list(map(int, input().split()))


def read_int():
    return int(input())


def read_str_list():
    return input().split()


def read_str():
    return input()


def solve():
    global c
    n = read_int()
    c = read_int_list()
    a, b = [], []
    if n > 1:
        a, b = map(list, zip(*[read_int_list() for i in range(n - 1)]))
    a = [ x - 1 for x in a]
    b = [ x - 1 for x in b]
    c = [ x - 1 for x in c]

    neighbors = defaultdict(list)
    for i in range(n - 1):
        neighbors[a[i]].append(b[i])
        neighbors[b[i]].append(a[i])

    parent = {}
    children = defaultdict(list)
    n_descendents = defaultdict(int)
    def dfs(i, p=-1):
        parent[i] = p
        n_descendents[i] = 0
        for neighbor in neighbors[i]:
            if neighbor == p:
                continue
            children[i].append(neighbor)
            dfs(neighbor, p=i)
            n_descendents[i] += 1 + n_descendents[neighbor]


    root = 0
    dfs(root, -1)

    # print('parent:', parent)
    # print('children:', children)
    # print('n_descendents:', n_descendents)

    results = []
    for k in range(n):

        # print()
        # print('k:', k)

        n0 = defaultdict(int)
        n1 = defaultdict(int)
        n2 = defaultdict(int)
        def count(i):


            # print('count, i:', i)

            global c
            for child in children[i]:
                count(child)
            n0[i] = int(c[i] == k)

            # n1
            if c[i] == k:
                n1[i] = n_descendents[i]
            else:
                for child in children[i]:
                    n1[i] += n0[child]
                    n1[i] += n1[child]

            # n2
            if c[i] == k:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += (1 + n_descendents[j]) * (1 + n_descendents[l])
            else:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += (1 + n_descendents[j]) * (1 + n_descendents[l])
                        n2[i] -= (1 + n_descendents[j] - n1[j] - n0[j]) * (1 + n_descendents[l] - n1[l] - n0[l])

        count(root)

        # print('n0:', n0)
        # print('n1:', n1)
        # print('n2:', n2)

        res = 0
        for i in range(n):
            res += n0[i] + n1[i] + n2[i]
        results.append(res)

        # print('res:', res)

    return results


def main():
    res = solve()
    print(*res, sep='\n')


if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys
import socket
from collections import defaultdict

hostnames = ['N551J', 'F551C', 'X553M']
input_file = 'f1.in'
if socket.gethostname() in hostnames:
    sys.stdin = open(input_file)


def read_int_list():
    return list(map(int, input().split()))


def read_int():
    return int(input())


def read_str_list():
    return input().split()


def read_str():
    return input()


def solve():
    global c
    n = read_int()
    c = read_int_list()
    a, b = map(list, zip(*[read_int_list() for i in range(n - 1)]))
    a = [ x - 1 for x in a]
    b = [ x - 1 for x in b]
    c = [ x - 1 for x in c]

    neighbors = defaultdict(list)
    for i in range(n - 1):
        neighbors[a[i]].append(b[i])
        neighbors[b[i]].append(a[i])

    parent = {}
    children = defaultdict(list)
    n_descendents = defaultdict(int)
    def dfs(i, p=-1):
        parent[i] = p
        n_descendents[i] = 0
        for neighbor in neighbors[i]:
            if neighbor == p:
                continue
            children[i].append(neighbor)
            dfs(neighbor, p=i)
            n_descendents[i] += 1 + n_descendents[neighbor]


    root = 0
    dfs(root, -1)

    # print('parent:', parent)
    # print('children:', children)
    # print('n_descendents:', n_descendents)

    results = []
    for k in range(n):

        # print()
        # print('k:', k)

        n0 = defaultdict(int)
        n1 = defaultdict(int)
        n2 = defaultdict(int)
        def count(i):


            # print('count, i:', i)

            global c
            for child in children[i]:
                count(child)
            n0[i] = int(c[i] == k)

            # n1
            if c[i] == k:
                n1[i] = n_descendents[i]
            else:
                for child in children[i]:
                    n1[i] += n0[child]
                    n1[i] += n1[child]

            # n2
            if c[i] == k:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += n_descendents[j] * n_descendents[l]
            else:
                for jj in range(len(children[i])):
                    for ll in range(jj):
                        j = children[i][jj]
                        l = children[i][ll]
                        n2[i] += n_descendents[j] * n_descendents[l]
                        n2[i] -= (n_descendents[j] - n1[j] - n0[j]) * (n_descendents[l] - n1[l] - n0[l])

        count(root)

        # print('n0:', n0)
        # print('n1:', n1)
        # print('n2:', n2)

        res = 0
        for i in range(n):
            res += n0[i] + n1[i] + n2[i]
        results.append(res)

        # print('res:', res)

    return results


def main():
    res = solve()
    print(*res, sep='\n')


if __name__ == '__main__':
    main()

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 -= 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)
        
        if reached[nxt]:
            continue
        size[p] = 0
        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)