Problem:


Score : 600 points

Problem Statement

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

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

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

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

Constraints

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

1 \leq c_i \leq N

1 \leq a_i,b_i \leq N

The given graph is a tree.

All values in input are integers.

Input

Input is given from Standard Input in the following format:

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

Output

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

Sample Input 1

3
1 2 1
1 2
2 3

Sample Output 1

5
4
0

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

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

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

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

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

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

Sample Input 2

1
1

Sample Output 2

1

Sample Input 3

2
1 2
1 2

Sample Output 3

2
2

Sample Input 4

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

Sample Output 4

5
8
10
5
5

Sample Input 5

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

Sample Output 5

18
15
0
14
23
0
23
0



Sample input:
3
1 2 1
1 2
2 3


Sample output:
5
4
0

Submission status: Accepted

Solution code:
def f(l):
    r = sum(l)**2
    for x in l:
        r -= x**2
    r = r//2
    r += sum(l)+1
    return r

N, = map(int, input().split())
G = [set() for _ in range(N+1)]
G2 = [set() for _ in range(N+1)]
sz = [1 for _ in range(N+1)]
X = [-1] + list(map(int, input().split()))
for _ in range(N-1):
    a, b = map(int, input().split())
    G[a].add(b)
    G[b].add(a)

stack = [1]
vs = set([1])
vs_dfs = []
while stack:
    v = stack.pop()
    vs_dfs.append(v)
    for u in G[v]:
        if u in vs:
            continue
        vs.add(u)
        G2[v].add(u)
        stack += [u]
G = None

for x in vs_dfs[::-1]:
    for u in G2[x]:
        sz[x] += sz[u]

stack = [1]
path = []
R = [0] * (N+1)
kh = [[] for _ in range(N+1)]
tt = [N-1 for _ in range(N+1)]
while stack:
    v = stack.pop()
    if v > 0:
        if path:
            parent = path[-1]
            kh[X[parent]].append(sz[v])
        path.append(v)
        tmp = []
        for u in G2[v]:
            tmp.append(sz[u])
            stack += [-v, u]
        if kh[X[v]]:
            tmp.append(kh[X[v]][-1]-1-sum(tmp))
            kh[X[v]][-1] -= sz[v]
        else:
            tmp.append(tt[X[v]]-sum(tmp))
            tt[X[v]] -= sz[v]
        R[X[v]] += f(tmp)
#        kz[X[v]].append(v)
    else:
#        kz[X[path[-1]]].pop()
        c = path.pop()
        kh[X[-v]].pop()
for i in range(1, N+1):
    print(R[i])

Submission status: Accepted

Solution code:
import bisect
import collections
import sys

sys.setrecursionlimit(1000000)
input = sys.stdin.readline
ACMOD = 1000000007
INF = 1 << 62


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


def llmi(n):
    return [lmi() for _ in range(n)]


N = int(input())
c_raw = lmi()
ab = llmi(N - 1)

adj = collections.defaultdict(list)
for a, b in ab:
    adj[a - 1].append(b - 1)
    adj[b - 1].append(a - 1)

total = N * (N + 1) // 2

visited = [False] * N

visited[0] = True
stacks = [[[0, 0]] for i in range(N)]
ans = [0] * N
curr = 1

# dfs = [0]
# while dfs:
#     node = dfs[-1]
#     prev, rem = stacks[c[node]][-1]
#     diff = curr - prev - rem
#     ans[c[node]] += (diff * (diff + 1)) // 2
#     stacks[c[node]][-1][1] += diff
#     if adj[node]:
#         next_node = adj[node].pop()
#         if not visited[next_node]:
#             visited[next_node] = True


color_stack = [[0, N] for _ in range(N)]
# color_groups = [[] for _ in range(N)]


class Tree:
    def __init__(self, i, _c):
        self.i = i
        self.color = _c
        self.child_count = 0
        self.children = []

    def set_child_count(self, parent=-1):
        for node_index in adj[self.i]:
            if node_index == parent:
                continue
            self.children.append(node_index)
            trees[node_index].set_child_count(self.i)
            self.child_count += trees[node_index].child_count + 1

    def dfs(self):
        color_stack[self.color][-1] -= self.child_count + 1
        for node_index in self.children:
            color_stack[self.color].append(trees[node_index].child_count + 1)
            trees[node_index].dfs()
            ans[self.color] -= get_path_count(color_stack[self.color].pop())
            # color_groups[self.color].append(color_stack[self.color].pop())


trees = [Tree(i, c - 1) for i, c in enumerate(c_raw)]


def get_path_count(x):
    return x + x * (x - 1) // 2


ans = [get_path_count(N)] * N
trees[0].set_child_count()
trees[0].dfs()

for color in range(N):
    # v = color_groups[color] + [color_stack[color][-1]]
    # print(N + N * (N - 1) // 2 - sum(map(lambda x: x + x * (x - 1) // 2, )))
    print(ans[color]-get_path_count(color_stack[color][-1]))

Submission status: Accepted

Solution code:
import bisect
import collections
import sys

sys.setrecursionlimit(1000000)
input = sys.stdin.readline
ACMOD = 1000000007
INF = 1 << 62


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


def llmi(n):
    return [lmi() for _ in range(n)]


N = int(input())
c_raw = lmi()
ab = llmi(N - 1)

adj = collections.defaultdict(list)
for a, b in ab:
    adj[a - 1].append(b - 1)
    adj[b - 1].append(a - 1)

total = N * (N + 1) // 2

visited = [False] * N

visited[0] = True
stacks = [[[0, 0]] for i in range(N)]
ans = [0] * N
curr = 1

# dfs = [0]
# while dfs:
#     node = dfs[-1]
#     prev, rem = stacks[c[node]][-1]
#     diff = curr - prev - rem
#     ans[c[node]] += (diff * (diff + 1)) // 2
#     stacks[c[node]][-1][1] += diff
#     if adj[node]:
#         next_node = adj[node].pop()
#         if not visited[next_node]:
#             visited[next_node] = True


color_stack = [[0, N] for _ in range(N)]
color_groups = [[] for _ in range(N)]


class Tree:
    def __init__(self, i, _c):
        self.i = i
        self.color = _c
        self.child_count = 0
        self.children = []

    def set_child_count(self, parent=-1):
        for node_index in adj[self.i]:
            if node_index == parent:
                continue
            self.children.append(node_index)
            trees[node_index].set_child_count(self.i)
            self.child_count += trees[node_index].child_count + 1

    def dfs(self):
        color_stack[self.color][-1] -= self.child_count + 1
        for node_index in self.children:
            color_stack[self.color].append(trees[node_index].child_count + 1)
            trees[node_index].dfs()
            color_groups[self.color].append(color_stack[self.color].pop())


trees = [Tree(i, c - 1) for i, c in enumerate(c_raw)]
trees[0].set_child_count()
trees[0].dfs()

for color in range(N):
    v = color_groups[color] + [color_stack[color][-1]]
    print(N + N * (N - 1) // 2 - sum(map(lambda x: x + x * (x - 1) // 2, v)))

Submission status: Accepted

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
# @pysnooper.snoop()
#import pysnooper # debug

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
    ROLLBACK = 1

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

    def add(self, pos, val):
        assert pos > 0, pos
        if self.status == BIT.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 sum_section(self, frm, to):
        to_val = self.sum(to)
        frm_val = self.sum(frm)
        return to_val - frm_val

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

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

class Graph:

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

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

def make_mybit(size):
    bit = BIT(size)
    for i in range(size):
        bit.add(i + 1, 1)
    bit.save()
    return bit

def init_color_s(size):
    color_s = [None] * graph.size
    for i in range(graph.size):
        color_s[i] = []
    return color_s

def dfs(graph):
    parent_s = [0] * graph.size
    in_s = [0] * graph.size
    out_s = [0] * graph.size
    order = 0 # starts from 1
    color_s = init_color_s(graph.size)
    def _dfs(v, prev):
        nonlocal order
        order += 1
        parent_s[v] = prev
        in_s[v] = order
        c = graph.vertices[v]
        color_s[c].append(v)
        for to in graph.edges[v]:
            if to == prev:
                continue
            _dfs(to, v)
        out_s[v] = order
    _dfs(0, -1)
    return parent_s, in_s, out_s, color_s

def calc(v):
    return v * (v + 1) // 2

def get_connected(bit, v, in_s, out_s):
    #print('v') # debug
    #print(v) # debug
    large_key = out_s[v]
    small_key = in_s[v] - 1
    return bit.sum_section(small_key, large_key)

def calc_for(bit, v, in_s, out_s):
    connected = get_connected(bit, v, in_s, out_s)
    #print('connected') # debug
    #print(connected) # debug
    return calc(connected)

def cut(bit, v, in_s, out_s):
    connected = get_connected(bit, v, in_s, out_s)
    bit.add(in_s[v], -1 * connected)

def solve(n, graph):
    bit = make_mybit(graph.size)
    parent_s, in_s, out_s, color_s = dfs(graph)
    #print('parent_s, in_s, out_s ') # debug
    #print(parent_s, in_s, out_s ) # debug
    #print('color_s') # debug
    #print(color_s) # debug
    for same_color_s in color_s:
        ans = calc(graph.size)
        bit.rollback()
        #print('bit.data') # debug
        #print(bit.data) # debug
        for v in same_color_s[::-1]:
            for to in graph.edges[v]:
                if to == parent_s[v]:
                    continue
                ans -= calc_for(bit, to, in_s, out_s)
            cut(bit, v, in_s, out_s)
        ans -= calc_for(bit, 0, in_s, out_s)
        #print('ans') # debug
        print(ans)






if __name__ == '__main__':
    n = int(input())
    graph = Graph(n)
    vertices = list(map(int, input().split()))
    for i, c in enumerate(vertices):
        vertices[i] = c - 1
    graph.vertices = vertices
    for _ in range(n - 1):
        frm, to = list(map(int, input().split()))
        frm -= 1
        to -= 1
        graph.add_edge(frm, to)
    solve(n, graph)


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

Submission status: Accepted

Solution code:
import sys
sys.setrecursionlimit(10**6)
 
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)
    
def dfs(u, p = -1):
 
    color = C[u]  # 色番号
    VC[color].append(u)  # 色ごとに stack しておく
    
    s = 1  # 部分木サイズ（を格納する）
    for v in E[u]:
        if v == p:
            continue
            
        child[u] = 0  # v の方向にある子孫のうち辿れる頂点の数 = s - 子孫にある、同色を根とする部分木のサイズの合計
        ret = dfs(v, u)
 
        s += ret  # 純粋な部分木のサイズを計算
        child[u] += ret  # 同色の頂点が切り離されたあとの部分木サイズ
        
        ans[color] -= child[u] * (child[u] + 1) // 2        
 
    VC[color].pop()
    if VC[color]: # 上位に同じ色がある場合
        child[VC[color][-1]] -= s  # 一番近い同色から、自分の部分木の大きさをひく
    else:  # 上位に同じ色がない場合
        root_size[color] -= s  # 上位に残っている頂点の数
        
    return s

child = [0] * N  # 辿っている子の中でのの残り数　とりあえず 0で初期化
root_size = [N] * N  # 根方向に残っている頂点の数（を格納する。操作前は N 全部）
VC = [ [] for _ in range(N)]  # 色番号ごと
ans = [N*(N+1)//2] * N  # 色ごとの（残りの）組み合わせ数　初期値は 全点の組み合わせ
 
dfs(0)
 
for i in range(N):
    print(ans[i] - root_size[i] * (root_size[i] + 1) // 2)

Submission status: Accepted

Solution code:
import sys
from operator import add
readline = sys.stdin.readline

def euler_tour(root, Edge):
    N = len(Edge)
    stack = [root]
    used = [False]*N
    tour = []
    down = []
    dist = [0]*N
    ta = tour.append
    sa = stack.append
    St = [None]*N
    En = [None]*N
    cnt = -1
    while stack:
        cnt += 1
        vn = stack.pop()
        ta(vn)
        
        if St[vn] is None:
            St[vn] = cnt
        En[vn] = cnt
        if used[vn]:
            continue
        for vf in Edge[vn]:
            if used[vf]:
                sa(vf)
            else:
                dist[vf] = 1+dist[vn]
                down.append(vf)
        stack.extend(down)
        down = []
        used[vn] = True
    return St, En, tour, dist

def parorder(Edge, p):
    N = len(Edge)
    par = [0]*N
    par[p] = -1
    stack = [p]
    order = []
    visited = set([p])
    ast = stack.append
    apo = order.append
    while stack:
        vn = stack.pop()
        apo(vn)
        for vf in Edge[vn]:
            if vf in visited:
                continue
            visited.add(vf)
            par[vf] = vn
            ast(vf)
    return par, order

def getcld(p):
    res = [[] for _ in range(len(p))]
    for i, v in enumerate(p):
        if v == -1:
            continue
        res[v].append(i)
    return res

class Segtree:
    def __init__(self, A, intv, initialize = True, segf = max):
        self.N = len(A)
        self.N0 = 2**(self.N-1).bit_length()
        self.intv = intv
        self.segf = segf
        if initialize:
            self.data = [intv]*self.N0 + A + [intv]*(self.N0 - self.N)
            for i in range(self.N0-1, 0, -1):
                self.data[i] = self.segf(self.data[2*i], self.data[2*i+1]) 
        else:
            self.data = [intv]*(2*self.N0)
            
    def __repr__(self):
        return str(' '.join(map(str, self.data[self.N0:self.N0+self.N])))

    
    def update(self, k, x):
        k += self.N0
        self.data[k] = x
        while k > 0 :
            k = k >> 1
            self.data[k] = self.segf(self.data[2*k], self.data[2*k+1])
    
    def query(self, l, r):
        L, R = l+self.N0, r+self.N0
        s = self.intv
        while L < R:
            if R & 1:
                R -= 1
                s = self.segf(s, self.data[R])
            if L & 1:
                s = self.segf(s, self.data[L])
                L += 1
            L >>= 1
            R >>= 1
        return s
    
    def binsearch(self, l, r, check, reverse = False):
        L, R = l+self.N0, r+self.N0
        SL, SR = [], []
        while L < R:
            if R & 1:
                R -= 1
                SR.append(R)
            if L & 1:
                SL.append(L)
                L += 1
            L >>= 1
            R >>= 1
        
        if reverse:
            for idx in (SR + SL[::-1]):
                if check(self.data[idx]):
                    break
            else:
                return -1
            while idx < self.N0:
                if check(self.data[2*idx+1]):
                    idx = 2*idx + 1
                else:
                    idx = 2*idx
            return idx - self.N0
        else:
            for idx in (SL + SR[::-1]):
                if check(self.data[idx]):
                    break
            else:
                return -1
            while idx < self.N0:
                if check(self.data[2*idx]):
                    idx = 2*idx
                else:
                    idx = 2*idx + 1
            return idx - self.N0

N = int(readline())
Col = list(map(lambda x: int(x)-1, readline().split()))
Colidx = [[] for _ in range(N)]
for i in range(N):
    Colidx[Col[i]].append(i)

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


root = 0
for root in range(N):
    if len(Edge[root]) == 1:
        break
P, L = parorder(Edge, root)
C = getcld(P)

St, En, tour, dist = euler_tour(root, Edge)

L = len(tour)
table = [0]*L
Ans = [0]*N
for i in range(N):
    table[En[i]] = 1

T = Segtree(table, 0, initialize = True, segf = add)
for c in range(N):
    Ci = Colidx[c]
    if Ci:
        Ci.sort(key = lambda x: dist[x], reverse = True)
        res = 0
        for v in Ci:
            sv = 0
            for cv in C[v]:
                k = T.query(St[cv], En[cv]+1)
                sv += k
                res += k*(k+1)//2
            T.update(En[v], -sv)
        k = T.query(St[root], En[root]+1)
        
        res += k*(k+1)//2
        for v in Ci:
            T.update(En[v], 1)
        Ans[c] = N*(N+1)//2-res
print('\n'.join(map(str, Ans)))

Submission status: Accepted

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

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


class BIT():
    def __init__(self, n, x=1):
        self.n = n
        self.T = [0]*(n+1)

    def add(self, i, x):
        i += 1
        while i <= self.n:
            self.T[i] += x
            i += i & -i

    def sum(self, i):
        i += 1
        ret = 0
        while i > 0:
            ret += self.T[i]
            i -= i & -i
        return ret

    def sump(self, l, r):
        return self.sum(r-1)-self.sum(l-1)


N = int(input())

C = [[] for _ in [0]*N]
for i, c in enumerate(map(int, input().split())):
    C[c-1].append(i)

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


pin = [-1]*N
pout = [-1]*N
k = 0


def dfs(v=0, p=-1):
    global k
    pin[v] = k
    k += 1
    for u in G[v]:
        if u == p:
            continue
        dfs(u, v)
    pout[v] = k


def f(x):
    return x*(x+1)//2


dfs()

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


total = f(N)
for c in C:
    c.sort(key=lambda x: pin[x], reverse=True)
    history = []
    ans = total
    for v in c:

        sub = 1
        for u in G[v]:
            if pin[u] < pin[v]:
                continue
            num = b.sump(pin[u], pout[u])
            ans -= f(num)
            sub += num

        b.add(pin[v], -sub)
        history.append((pin[v], sub))

    ans -= f(b.sump(0, N))
    for h in history:
        b.add(*h)
    print(ans)

Submission status: Accepted

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

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)

Submission status: Accepted

Solution code:
import sys
import itertools

sys.setrecursionlimit(3 * 10**5)


class Node:
  def __init__(self, color):
    self.color = int(color) - 1
    self.children = []

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

  @staticmethod
  def prepare_nodes(n):
    nodes = tuple(map(Node, 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.color
    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))

  for i in range(n):
    update_ans(s, i)
    print(ans_list[i])


main()

Submission status: Accepted

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

Submission status: Accepted

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


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

# 回答用
ans = [0]*N


def f(curr, par=-1):
    # 再帰関数
    # curr: 現在の節点
    # par : 親節点の番号
    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[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: Accepted

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[1:], sep="\n")

Submission status: Accepted

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):
    print(total - ans[i])

Submission status: Accepted

Solution code:

import sys
# 小さすぎても怒られるので
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):
    print(total - ans[i])

Submission status: Accepted

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=[0] # これが1からツアーで辿った点
USED=[0]*(N+1)
Parent=[-1]*(N+1)
TOP_SORT=[1]
Children=[[] for i in range(N+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
            Children[x].append(to)
            TOP_SORT.append(to)
        else:
            QUE.append(to)
    QUE.extend(QUE2)
    QUE2=deque()

    USED[x]=1

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

for x in TOP_SORT:
    CTOP_SORT[C[x]].append(x)

EFIRST=[-1]*(N+1)
ELAST=[-1]*(N+1)
for i in range(len(EULER)):
    if EFIRST[EULER[i]]==-1:
        EFIRST[EULER[i]]=i
    ELAST[EULER[i]]=i
        
    CE[C[EULER[i]]].append(i)

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

Size[0]=N+1


LEN=len(EULER)

BIT=[0]*(LEN+1) # 1-indexedなtree

def update(v,w): # index vにwを加える
    while v<=LEN:
        BIT[v]+=w
        v+=(v&(-v)) # 自分を含む大きなノードへ. たとえばv=3→v=4

def getvalue(v): # [1,v]の区間の和を求める
    ANS=0
    while v!=0:
        ANS+=BIT[v]
        v-=(v&(-v)) # 自分より小さい2ベキのノードへ. たとえばv=3→v=2へ
    return ANS

for i in range(1,N+1):
    update(EFIRST[i],1)

for color in range(1,N+1):
    if CE[color]==[]:
        print(0)
        continue

    EULER2=CE[color]
    TOP_SORT2=CTOP_SORT[color]
    HISTORY=[]

    #print(EULER2,TOP_SORT2)
    #print(BIT)

    ANS=0
    for x in TOP_SORT2[::-1]:
        for to in Children[x]:
            sc=getvalue(ELAST[to])-getvalue(EFIRST[to]-1)
            update(EFIRST[x],-sc)
            HISTORY.append((EFIRST[x],sc))
            ANS+=sc*(sc+1)//2
        update(EFIRST[x],-1)
        HISTORY.append((EFIRST[x],1))

    #print(ANS)

    sc=getvalue(len(EULER))
    ANS+=sc*(sc+1)//2

    print(N*(N+1)//2-ANS)
    for x,y in HISTORY:
        update(x,y)