human
Solve in Python:
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-----
3
1 2 1
1 2
2 3

-----Sample Output-----
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.
 Write me a Python program to find the number of simple paths that visit a vertex painted in the color k one or more times?
assistant
#Here is Python program to find the number of simple paths that visit a vertex painted in the color k one or more times.

#!/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)
        self.E[t].append(f)

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 e:
        d[k] += e[k]
    return d

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] = list(map(int, input().split()))

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

ans = [0]*N

ret = {}
last_order, parent, size = make_order(g, 0)
for curr in last_order:
    cn = c[curr]
    rrr = defaultdict(int)
    for dest in g.E[curr]:
        if dest == parent[curr]:
            continue
        child = ret.pop(dest)

        n = size[dest]-child[cn]
        ans[cn] += n*(n+1)//2

        # マージ
        rrr = merge(rrr, child)

    rrr[cn] = size[curr]
    ret[curr] = rrr

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