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: Time Limit Exceeded

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

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

    def add(self, pos, val):
        assert(pos > 0)
        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

class Graph:

    def __init__(self, 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] = set()

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

    def size(self):
        return len(self.vertices) - 1

    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
        base_bit = self.make_base_bit()
        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
            bit = base_bit.make_copy()
            complement = self.solve(colored_v_s, bit)
            #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: Time Limit Exceeded

Solution code:
import sys
sys.setrecursionlimit(10**8)
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: Time Limit Exceeded

Solution code:
#!/usr/bin/env python3

from numpy import ones, zeros

N = int(input())

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

class Leaf:
    def __init__(self, color):
        self.root = None
        self.child = []
        self.color = color
        self.ret = None

    def add_leaf(self, aleaf):
        self.child.append(aleaf)
        aleaf.root = self

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

leaf_set = [Leaf(_) for _ in c]
stack = [0]
while stack:
    #print(stack)
    a = stack.pop()
    for b in link[a]:
        if leaf_set[b] == leaf_set[a].root:
            continue
        leaf_set[a].add_leaf(leaf_set[b])
        stack.append(b)

#for i in leaf_set:
#    print(i, i.root, i.child)

root = leaf_set[0]
ret = zeros(N).astype(int)
stack = [root]
while stack:
    aleaf = stack.pop()
    if aleaf.child and aleaf.child[0].ret is None:
        stack.append(aleaf)
        stack += aleaf.child
        continue
    aleaf.ret = ones(N).astype(int)
    for cld in aleaf.child:
        n = cld.ret[aleaf.color]
        ret[aleaf.color] += n * (n + 1) // 2
        aleaf.ret += cld.ret
    aleaf.ret[aleaf.color] = 0

n = root.ret
ret += n * (n + 1) // 2

#for i in leaf_set:
#    print(i.color, i.ret)

ret = N * (N + 1) // 2 - ret
for r in ret:
    print(r)

Submission status: Time Limit Exceeded

Solution code:
#n = 5
#C = list(map(int,"1 2 3 4 5".split()))
#line = ["1 2", "2 3", "3 4", "3 5"]
import sys
sys.setrecursionlimit(10**6)
input = sys.stdin.readline

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

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

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

reached[0] = True
dfs (0)

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

Submission status: Time Limit Exceeded

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:
    ret[curr] = defaultdict(int)
    for dest in g.E[curr]:
        if dest == parent[curr]:
            continue
        child = ret[dest]

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

        # マージ
        ret[curr] = merge(ret[curr], child)

    ret[curr][c[curr]] = size[curr]


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

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

Submission status: Time Limit Exceeded

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

# 回答用
ans = [0]*N


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

        # 自身の色と同じ場合、子の頂点の数から加算
        n = 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]] = 0

    for color in range(N):
        if color != c[curr]:
            ret[color] += 1
    return ret


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

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

Submission status: Time Limit Exceeded

Solution code:
import sys
import itertools

sys.setrecursionlimit(3 * 10**5)


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

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

  def dfs(self, prev_id):
    result = {self.id}
    for node in self.relations:
      if node.id != prev_id:
        result |= node.dfs(self.id)

    return result


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

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

  return node_dict


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

  paths_dict = {k: set() for k in range(1, n + 1)}

  for node in nodes:
    id_sets = [relation.dfs(node.id) for relation in node.relations]

    paths = {(node.id, node.id)}
    for id_set in id_sets:
      paths |= {tuple(sorted((node.id, id))) for id in id_set}

    cmbs = itertools.combinations(range(len(id_sets)), 2)
    for i, j in cmbs:
      paths |= {
        tuple(sorted((a, b)))
        for a in id_sets[i]
        for b in id_sets[j]
      }

    paths_dict[node.c] |= paths


  for paths in paths_dict.values():
    print(len(paths))


main()

Submission status: Time Limit Exceeded

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

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

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

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

Submission status: Time Limit Exceeded

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

Submission status: Time Limit Exceeded

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

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

# 回答用
ans = [0]*N


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

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


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

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

Submission status: Time Limit Exceeded

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

Submission status: Time Limit Exceeded

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

Submission status: Time Limit Exceeded

Solution code:
import sys
sr = lambda: sys.stdin.readline().rstrip()
ir = lambda: int(sr())
lr = lambda: list(map(int, sr().split()))
def c2(d):
    return d*(d+1)//2
def resolve():
    N = ir()
    c = lr()
    ab = [[] for i in range(N)]
    for i in range(N-1):
        a, b = lr()
        ab[a-1].append(b-1)
        ab[b-1].append(a-1)
    tot = c2(N)
    for i in range(N):
        cnt = [0]*N
        vis = [0]*N
        s = [[0, 0]]
        vis[0] += 1
        while len(s) > 0:
            v, x = s.pop()
            if c[v] != i+1:
                cnt[x] += 1
            for j in ab[v]:
                if vis[j] == 0:
                    if c[v] == i+1:
                        x += 1
                    s.append([j, x])
                    vis[j] += 1
        print(tot-sum(map(c2, cnt)))
resolve()

Submission status: Time Limit Exceeded

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)

def dfs(ns, prev, now, toru, results):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    tmp_s2_sum = 0
    usList = defaultdict(lambda: [0,0,0]) # sum(v), sum(v*s), sum(v**2)
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = results[u]  # dfs(ns, now, u, toru)
        for kk,vv in tmp_end.items():
            if kk!=k:
                end[kk] += vv
                usList[kk][0] += vv # .append((vv, tmp_s))
                usList[kk][1] += (vv*tmp_s)
                usList[kk][2] += vv**2
        end[k] += tmp_s
        s += tmp_s
        tmp_s2_sum += tmp_s**2
    for kk,vv in usList.items():
        if kk!=k:
            toru[kk] += (s - 1) * vv[0] - vv[1] - (vv[0]**2 - vv[2]) / 2
            toru[kk] += vv[0]
#             x_sum = 0
#             x2_sum = 0
#             for item in vv:
#                 toru[kk] += (s - item[1]) * item[0]
#                 x_sum += item[0]
#                 x2_sum += item[0]**2
#             toru[kk] -= (x_sum**2 - x2_sum) / 2
    toru[k] += (1 + ((s-1)**2 - tmp_s2_sum)/2 + (s-1))
#     print(now, prev, end, toru, s, usList)
    results[now] = end, toru, s

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

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

Submission status: Time Limit Exceeded

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)

def dfs(ns, prev, now, toru, results):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    tmp_s2_sum = 0
    usList = defaultdict(list)
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = results[u]  # dfs(ns, now, u, toru)
        for kk,vv in tmp_end.items():
            if kk!=k:
                end[kk] += vv
            usList[kk].append((vv, tmp_s))
        end[k] += tmp_s
        s += tmp_s
        tmp_s2_sum += tmp_s**2
    for kk,vv in usList.items():
        if kk!=k:
            x_sum = 0
            x2_sum = 0
            for item in vv:
                toru[kk] += (s - item[1]) * item[0]
                x_sum += item[0]
                x2_sum += item[0]**2
            toru[kk] -= (x_sum**2 - x2_sum) / 2
    toru[k] += (1 + ((s-1)**2 - tmp_s2_sum)/2 + (s-1))
#     print(now, prev, end, toru, s)
    results[now] = end, toru, s

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

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

Submission status: Time Limit Exceeded

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)

def dfs(ns, prev, now, toru, results):
    end = defaultdict(int)
    k = cs[now]
    end[k] = 1
    s = 1 # nowで終わるパスの種類
    tmp_s2_sum = 0
    usList = defaultdict(list)
    for u in ns[now]:
        if u==prev:
            continue
        tmp_end, tmp_toru, tmp_s = results[u]  # dfs(ns, now, u, toru)
        for kk,vv in tmp_end.items():
            if kk!=k:
                end[kk] += vv
            usList[kk].append((vv, tmp_s))
        end[k] += tmp_s
        s += tmp_s
        tmp_s2_sum += tmp_s**2
    for kk,vv in usList.items():
        if kk!=k:
            x_sum = 0
            x2_sum = 0
            for item in vv:
                toru[kk] += (s - item[1]) * item[0]
                x_sum += item[0]
                x2_sum += item[0]**2
            toru[kk] -= (x_sum**2 - x2_sum) / 2
    toru[k] += (1 + ((s-1)**2 - tmp_s2_sum)/2 + (s-1))
#     print(now, prev, end, toru, s)
    results[now] = end, toru, s

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

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

Submission status: Time Limit Exceeded

Solution code:
# 解説放送そのまま
from collections import defaultdict
import sys

sys.setrecursionlimit(10 ** 6)
int1 = lambda x: int(x) - 1
p2D = lambda x: print(*x, sep="\n")
def II(): return int(sys.stdin.readline())
def MI(): return map(int, sys.stdin.readline().split())
def MI1(): return map(int1, sys.stdin.readline().split())
def LI(): return list(map(int, sys.stdin.readline().split()))
def LI1(): return list(map(int1, sys.stdin.readline().split()))
def LLI(rows_number): return [LI() for _ in range(rows_number)]
def SI(): return sys.stdin.readline()[:-1]

class BitSum:
    def __init__(self, n):
        self.n = n + 3
        self.table = [0] * (self.n + 1)

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

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

def main():
    def dfs(u=0,pu=-1):
        ll[u]=len(uu)
        uu.append(u)
        for v in to[u]:
            if v==pu:continue
            dfs(v,u)
        rr[u]=len(uu)

    def sumlr(l,r):
        return bit.sum(r-1)-bit.sum(l-1)

    n=II()
    cc=LI1()
    to=[[] for _ in range(n)]
    for _ in range(n-1):
        a,b=MI1()
        to[a].append(b)
        to[b].append(a)
    ll=[-1]*n
    rr=[-1]*n
    uu=[]
    dfs()
    bit=BitSum(n)
    for u in range(n):bit.add(u,1)
    ctou=defaultdict(list)
    for u,c in enumerate(cc):ctou[c].append(u)
    pn=lambda x:x*(x+1)//2
    for c in range(n):
        log=[]
        ans = pn(n)
        for u in sorted(ctou[c],key=lambda x:-ll[x]):
            us=1
            for v in to[u]:
                if ll[v]<ll[u]:continue
                vs=sumlr(ll[v],rr[v])
                ans-=pn(vs)
                us+=vs
            bit.add(ll[u],-us)
            log.append((ll[u],us))
        vs = sumlr(0,n)
        ans -= pn(vs)
        print(ans)
        for i,s in log:bit.add(i,s)

main()