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:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np
from numba import njit


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


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

tour, ind_L, ind_R, parent = map(lambda x: np.array(x, np.int64), EulerTour(graph))


@njit('i8[:](i8[:],i8[:],i8[:],i8[:],i8[:])', cache=True)
def main(C, tour, ind_L, ind_R, parent):
    N = len(C) - 1
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


answer = main(C, tour, ind_L, ind_R, parent)
print('\n'.join(answer[1:].astype(str)))

Submission status: Accepted

Solution code:
import sys
from collections import defaultdict as dd
input = sys.stdin.readline
N = int(input())
c = [0] + list(map(int, input().split()))
e = dd(list)
for _ in range(N - 1):
  u, v = map(int, input().split())
  e[u].append(v)
  e[v].append(u)
s = [1]
order = []
vis = [0] * (N + 1)
ps = [0] * (N + 1)
while len(s):
  x = s.pop()
  vis[x] = 1
  order.append(x)
  for y in e[x]:
    if vis[y]: continue
    s.append(y)
    ps[y] = x

size = [1] * (N + 1)
res = [0] * (N + 1)
d = [dd(int) for _ in range(N + 1)]
for x in order[: : -1]:
  for y in e[x]:
    if y == ps[x]: continue
    size[x] += size[y]
    res[c[x]] += (size[y] - d[y][c[x]]) * (size[y] - d[y][c[x]] + 1) // 2

    if len(d[x]) < len(d[y]): d[x], d[y] = d[y], d[x]
    for z in d[y]: d[x][z] += d[y][z]
  d[x][c[x]] = size[x]


for i in range(1, N + 1):
  print(N * (N + 1) // 2 - res[i] - (N - d[1][i]) * (N - d[1][i] + 1) // 2)

Submission status: Accepted

Solution code:
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
import numpy as np
from numba import njit


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


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

tour, ind_L, ind_R, parent = map(lambda x: np.array(x, np.int64), EulerTour(graph))


@njit(cache=True)
def main(C, tour, ind_L, ind_R, parent):
    removed = np.zeros(N + 1, np.int64)
    answer = np.full_like(removed, N * (N + 1) // 2)
    memo = np.zeros_like(answer)
    for v in tour:
        if v > 0:
            memo[v] = removed[C[parent[v]]]
        else:
            v = -v
            removed[C[v]] += 1
            p = parent[v]
            x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
            x -= removed[C[p]] - memo[v]
            answer[C[p]] -= x * (x + 1) // 2
            removed[C[p]] += x
    for i, x in enumerate(removed):
        x = N - x
        answer[i] -= x * (x + 1) // 2
    return answer


answer = main(C, tour, ind_L, ind_R, parent)
print('\n'.join(answer[1:].astype(str)))

Submission status: Accepted

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


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


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

tour, ind_L, ind_R, parent = EulerTour(graph)
removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in tour:
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        c = C[parent[v]]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
        x -= removed[c] - memo[v]
        answer[c] -= x * (x + 1) // 2
        removed[c] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))

Submission status: Accepted

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


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


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

tour, ind_L, ind_R, parent = EulerTour(graph)
removed = [0] * (N + 1)
answer = [N * (N + 1) // 2] * (N + 1)
memo = [0] * (N + 1)

for v in tour:
    if v > 0:
        memo[v] = removed[C[parent[v]]]
    else:
        v = -v
        removed[C[v]] += 1
        p = parent[v]
        x = (ind_R[v] - ind_L[v]) // 2 + 1  # subtree size
        x -= removed[C[p]] - memo[v]
        answer[C[p]] -= x * (x + 1) // 2
        removed[C[p]] += x

for i, x in enumerate(removed):
    x = N - x
    answer[i] -= x * (x + 1) // 2
print('\n'.join(map(str, answer[1:])))

Submission status: Accepted

Solution code:
import sys
sys.setrecursionlimit(10**6)
def main():
    n=int(input())
    color=list(map(int,input().split()))
    nodes=[[]for i in range(n)]
    cut=[[0]for i in range(n+1)]
    ans=[0]*(n+1)
    for i in range(n-1):
        a,b=map(int,input().split())
        nodes[a-1].append(b-1)
        nodes[b-1].append(a-1)
    def num(p,parent):
        s=0
        for child in nodes[p]:
            if child==parent:
                continue
            cut[color[p]].append(0)
            nc=num(child,p)
            group=nc-cut[color[p]].pop()
            ans[color[p]]+=group*(group+1)//2
            s+=nc
        s+=1
        cut[color[p]][-1]+=s
        return s
    num(0,-1)
    for a,c in zip(ans[1:],cut[1:]):
        group=n-c[0]
        c=group*(group+1)//2
        print(n*(n+1)//2-c-a)
main()

Submission status: Accepted

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

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

ret = [N * (N - 1) // 2 + N] * N

par = [[N + i] for i in range(N)]
vec2 = [[] for _ in range(N * 2)]
have = [0] * (N * 2)
way = [0] * (N * 2)
fr = [0] * (N + 2)

for i in range(N, 2 * N) :
  have[i] = N + 1

mem = defaultdict(int)

def dfs(cur, pre) :
  s = 1
  ancestor = par[c[cur]][-1]
  vec2[ancestor].append(cur)
  if pre != -1 :
    fr[cur] = way[ancestor]
  par[c[cur]].append(cur)
  for nex in vec[cur] :
    if nex == pre : continue
    way[cur] = nex
    s += dfs(nex, cur)
    mem[(cur, nex)] = have[nex]
  par[c[cur]].pop()
  have[cur] = s
  return s
  
dfs(0, -1)

for i in range(N, 2 * N) :
  mem[(i, 0)] = N

for i in range(N) :
  q = [(N + i, -1)]
  while q :
    cur, pre = q.pop()
    for nex in vec2[cur] :
      mem[(cur, fr[nex])] -= have[nex]
      q.append((nex, cur))

for k, v in mem.items() :
  if v <= 0: continue
  if k[0] >= N :
    ret[k[0] % N] -= v * (v - 1) // 2 + v
  else :
    ret[c[k[0]]] -= v * (v - 1) // 2 + v
  
for i in range(N) :
  print(ret[i])

Submission status: Accepted

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


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

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

    def merge(A, B):
        if len(A) > len(B):
            A, B = B, A
        for a in A:
            if a in B:
                B[a] += A[a]
            else:
                B[a] = A[a]
        return B

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

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

    size = [1] * (N+1)
    dic = [{} for _ in range(N+1)]
    ans = [calc(N)] * (N+1)
    for v in seq:
        for u in child[v]:
            size[v] += size[u]
        c = C[v-1]
        for u in child[v]:
            if c in dic[u]:
                ans[c] -= calc(size[u] - dic[u][c])
            else:
                ans[c] -= calc(size[u])
            dic[v] = merge(dic[v], dic[u])
        dic[v][c] = size[v]
    for c in range(1, N+1):
        if c in dic[1]:
            print(ans[c] - calc(N - dic[1][c]))
        else:
            print(ans[c] - calc(N))


if __name__ == '__main__':
    main()

Submission status: Accepted

Solution code:
import sys
from collections import defaultdict

# 再帰制限を緩和するおまじない
sys.setrecursionlimit(10**6)

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

# 部分木のサイズと、「色iを封鎖したときに到達できない頂点の個数」を持つ辞書を返す
def dfs(v,p):
    ret=defaultdict(int)
    size=1
    for vv in g[v]:
        if vv==p:
            continue
        ss,d=dfs(vv,v)
        size+=ss
        ans[c[v]]+=sum(ss-d[c[v]])
        
        # マージテク
        if len(ret)<len(d):
            ret,d=d,ret
        for vvv in d:
            ret[vvv]+=d[vvv]
    ret[c[v]]=size
    return size,ret

n=int(input())
c=list(map(int,input().split()))
g=[[] for _ in range(n+1)]
ans=[0]*(n+1)
for _ in range(n-1):
    s,t=map(int, input().split())
    s-=1
    t-=1
    g[s].append(t)
    g[t].append(s)
_,ret=dfs(0,-1)
for i in range(1,n+1):
    print(sum(n)-ans[i]-sum(n-ret[i]))

Submission status: Accepted

Solution code:
from collections import defaultdict
import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

N = int(readline())
C = tuple(map(lambda x: int(x)-1, readline().split()))

temp = map(int,read().split())
temp = zip(temp,temp)

adj = [list() for _ in range(N)]
for a,b in temp:
    adj[a-1].append(b-1)
    adj[b-1].append(a-1)


parent = [None]*N
parent[0] = -1

stack = [0]
order = []
while stack:
    p = stack.pop()
    order.append(p)
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            stack.append(c)

def merge(a,b):
    if len(a) > len(b):
        a,b = b,a
    for k,v in a.items():
        b[k] += v
    return b

dp = [defaultdict(int) for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [N*(N+1)//2]*N

for c in reversed(order):
    p = parent[c]
    if p >= 0:
        dp[c][C[c]] = subtreesize[c]+1
        x = subtreesize[c]-dp[c][C[p]]+1
        cnt[C[p]] -= x*(x+1)//2
        subtreesize[p] += subtreesize[c]+1

        dp[p] = merge(dp[c],dp[p])
        dp[c] = None

dp[0][C[0]] = subtreesize[0]+1
for c in range(N):
    x = subtreesize[0]-dp[0][c]+1
    cnt[c] -= x*(x+1)//2

print(*cnt, sep='\n')

Submission status: Accepted

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

degs = [0] * N
clusters = [[] for _ in range(N)]


def go(curr, fro):
    col = Cs[curr] - 1
    deg = 1
    clust = {}
    for n in adj[curr]:
        if n == fro:
            continue
        ch_deg, ch_clust = yield go(n, curr)
        clusters[col].append(ch_deg - ch_clust.get(col, 0))
        if len(ch_clust) > len(clust):
            clust, ch_clust = ch_clust, clust
        for k, v in ch_clust.items():
            clust[k] = clust.get(k, 0) + v
        deg += ch_deg
    clust[col] = deg
    return deg, clust


def tramp(gen):
    stack = [gen]
    val = None
    while stack:
        try:
            stack.append(stack[-1].send(val))
            val = None
        except StopIteration as e:
            stack.pop()
            val = e.value
    return val


_, root_clust = tramp(go(0, -1))
for i in range(N):
    clusters[i].append(N - root_clust.get(i, 0))


for clust in clusters:
    tot = (N * N + N) // 2
    for cnt in clust:
        tot -= (cnt * cnt + cnt) // 2
    print(tot)

Submission status: Accepted

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

degs = [0] * N
clusters = [[] for _ in range(N)]


def go(curr, fro):
    col = Cs[curr] - 1
    deg = 1
    clust = {}
    for n in adj[curr]:
        if n == fro:
            continue
        ch_deg, ch_clust = yield n, curr
        clusters[col].append(ch_deg - ch_clust.get(col, 0))
        if len(ch_clust) > len(clust):
            clust, ch_clust = ch_clust, clust
        for k, v in ch_clust.items():
            clust[k] = clust.get(k, 0) + v
        deg += ch_deg
    clust[col] = deg
    return deg, clust


def tramp(gen_fn, *args):
    stack = [gen_fn(*args)]
    val = None
    while stack:
        gen = stack[-1]
        try:
            params = gen.send(val)
        except StopIteration as e:
            stack.pop()
            val = e.value
        else:
            stack.append(go(*params))
            val = None
    return val


_, root_clust = tramp(go, 0, -1)
for i in range(N):
    clusters[i].append(N - root_clust.get(i, 0))


for clust in clusters:
    tot = (N * N + N) // 2
    for cnt in clust:
        tot -= (cnt * cnt + cnt) // 2
    print(tot)

Submission status: Accepted

Solution code:
import sys
from collections import defaultdict

sys.setrecursionlimit(10 ** 6)


def dfs(v, p, ccc, links, ans):
    ret_colors = defaultdict(int)
    ret_count = 1
    cv = ccc[v]
    for u in links[v]:
        if u == p:
            continue
        sub_colors, sub_count = dfs(u, v, ccc, links, ans)
        cc = sub_count - sub_colors[cv]
        ans[cv] -= cc * (cc + 1) // 2
        ret_count += sub_count

        if len(ret_colors) < len(sub_colors):
            ret_colors, sub_colors = sub_colors, ret_colors

        for c, cnt in sub_colors.items():
            ret_colors[c] += cnt

    ret_colors[cv] = ret_count

    return ret_colors, ret_count


def solve(n, ccc, links):
    if n == 1:
        return [1]

    all_pair = n * (n + 1) // 2
    ans = [all_pair] * (n + 1)

    colors, count = dfs(0, -1, ccc, links, ans)
    assert count == n

    for c in range(1, n + 1):
        cc = n - colors[c]
        ans[c] -= cc * (cc + 1) // 2

    return ans[1:]


n, *cab = map(int, sys.stdin.buffer.read().split())
ccc = cab[:n]
links = [set() for _ in range(n)]

for a, b in zip(cab[n + 0::2], cab[n + 1::2]):
    a -= 1
    b -= 1
    links[a].add(b)
    links[b].add(a)

print('\n'.join(map(str, solve(n, ccc, links))))

Submission status: Accepted

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

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
AB = zip(m, m)

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)

root = 1
parent = [0] * (N + 1)
depth = [0] * (N + 1)  # 根を深さ 1 として
order = []
stack = [root]
depth[root] = 1
while stack:
    x = stack.pop()
    order.append(x)
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        depth[y] = depth[x] + 1


def merge(A, B):
    if len(A) < len(B):
        A, B = B, A
    for key, cnt in B.items():
        A[key] += cnt
    return A


dp = [defaultdict(int) for _ in range(N + 1)]  # 部分木内で、消去済の頂点数の分布
size = [0] * (N + 1)  # 部分木の大きさ
answer = [N * (N + 1) // 2] * (N + 1)

for v in order[::-1]:
    size[v] += 1
    p = parent[v]
    dp[v][C[v]] = size[v]
    if not p:
        break
    size[p] += size[v]
    c = C[p]
    x = size[v] - dp[v][c]
    answer[c] -= x * (x + 1) // 2
    dp[p] = merge(dp[p], dp[v])

A = dp[root]
for c in range(1, N + 1):
    x = N - A[c]
    answer[c] -= x * (x + 1) // 2

print(*answer[1:], sep='\n')

Submission status: Accepted

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

ns = lambda: stdin.readline().rstrip()
ni = lambda: int(stdin.readline().rstrip())
nm = lambda: map(int, stdin.readline().split())
nl = lambda: list(map(int, stdin.readline().split()))

n = ni()
color = nl()
g = [list() for _ in range(n)]
for _ in range(n-1):
    a, b = nm()
    g[a-1].append(b-1)
    g[b-1].append(a-1)

size = [1]*n
par = [-1]*n
h = [list() for _ in range(n+1)]
idx = [0]*n
def dfs(v, p):
    h[color[v]].append(v)
    for x in g[v]:
        if x == p: continue
        par[x] = v
        size[v] += dfs(x, v)
        h[color[v]].append(v)
    h[color[v]].append(v)
    return size[v]
dfs(0, -1)
# print(size)
for v in range(n):
    if idx[v] < len(g[v]) and g[v][idx[v]] == par[v]:
        idx[v] += 1
acc = [0]*(n+10)
ans = [n*(n+1)//2]*(n+1)
for col in range(1, n+1):
    # print(h[col])
    acc[-1] = 0
    p = [-1]
    for v in h[col]:
        if p[-1] == v:
            p.pop()
            if idx[v] >= len(g[v]):
                acc[p[-1]] += size[v]
                continue
            q = size[g[v][idx[v]]] - acc[v]
            idx[v] += 1
            if idx[v] < len(g[v]) and g[v][idx[v]] == par[v]:
              idx[v] += 1
            ans[col] -= q*(q+1)//2
            acc[v] = 0
            # print(v, q)
            p.append(v)
        else:
            p.append(v)
    q = n - acc[-1]
    # print(-1, q)
    ans[col] -= q*(q+1)//2
print(*ans[1:], sep='\n')

Submission status: Accepted

Solution code:
import collections
import sys

sys.setrecursionlimit(10 ** 6)

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

subtree_sizes = [0] * N


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


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


def calc_components(node, parent):
    color_stack_pushed = False
    if parent >= 0 and colors[parent] != colors[node]:
        color_parent = color_stack[colors[node]][-1]
        components[colors[node]][color_parent] -= subtree_sizes[node]
        components[colors[parent]][node] = subtree_sizes[node]
        color_stack[colors[parent]].append(node)
        color_stack_pushed = True
    for i in neighbors[node]:
        if i != parent:
            calc_components(i, node)
    if color_stack_pushed:
        color_stack[colors[parent]].pop()


calc_components(0, -1)


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


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

Submission status: Accepted

Solution code:
import sys
input = lambda: sys.stdin.readline().rstrip()
N = int(input())
C = [int(a) - 1 for a in input().split()]
X = [[] for i in range(N)]
for i in range(N-1):
    x, y = map(int, input().split())
    X[x-1].append(y-1)
    X[y-1].append(x-1)

def EulerTour(n, X, i0):
    f = lambda k: k * (k + 1) // 2
    USED = [0] * n
    ORG = [0] * n
    TMP = [0] * n
    P = [-1] * n
    Q = [~i0, i0]
    ct = -1
    ET1 = [0] * n
    ET2 = [0] * n
    ANS = [f(n)] * n
    while Q:
        i = Q.pop()
        if i < 0:
            i = ~i
            ET2[i] = ct
            USED[C[i]] += 1 + TMP[i]
            if i:
                p = P[i]
                k = ET2[i] - ET1[i] + 1 - USED[C[p]] + ORG[i]
                ANS[C[p]] -= f(k)
                TMP[p] += k
            continue
        if i >= 0:
            if i: ORG[i] = USED[C[P[i]]]
            ct += 1
            if ET1[i] == 0: ET1[i] = ct
        for a in X[i][::-1]:
            if a != P[i]:
                P[a] = i
                for k in range(len(X[a])):
                    if X[a][k] == i:
                        del X[a][k]
                        break
                Q.append(~a)
                Q.append(a)
    for i in range(n):
        ANS[i] -= f(n - USED[i])
    return ANS

print(*EulerTour(N, X, 0), sep = "\n")

Submission status: Accepted

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

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
AB = zip(m, m)

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)


root = 1
parent = [0] * (N + 1)
depth = [0] * (N + 1)  # 根を深さ 1 として
order = []
stack = [root]
depth[root] = 1
while stack:
    x = stack.pop()
    order.append(x)
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        depth[y] = depth[x] + 1


def merge(v, p):
    A = dp[v]
    B = dp[p]
    size_A = size[v]
    size_B = size[p]
    c = C[p]
    if len(A) < len(B):
        A, B = B, A
        size_A, size_B = size_B, size_A
    answer[c] += (size_A - A[c]) * (size_B - B[c])
    # B を A にマージ
    for key, cnt in B.items():
        answer[key] -= cnt * A[key]
        A[key] += cnt
    size[p] = size_A + size_B
    dp[p] = A


answer = [0] * (N + 1)
dp = [defaultdict(int) for _ in range(N + 1)]
size = [0] * (N + 1)

for v in order[::-1]:
    c = C[v]
    n = size[v] + 1 - dp[v][c]
    answer[c] -= size[v] * n
    size[v] += 1
    dp[v][c] = size[v]
    if v == root:
        break
    p = parent[v]
    merge(v, p)

A = dp[root]
c = size[root]
for v in range(1, N + 1):
    print(c * A[v] + answer[v])

Submission status: Accepted

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

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import defaultdict

N = int(readline())
C = (0,) + tuple(map(int, readline().split()))
m = map(int, read().split())
AB = zip(m, m)

graph = [[] for _ in range(N + 1)]
for a, b in AB:
    graph[a].append(b)
    graph[b].append(a)


root = 1
parent = [0] * (N + 1)
depth = [0] * (N + 1)  # 根を深さ 1 として
order = []
stack = [root]
depth[root] = 1
while stack:
    x = stack.pop()
    order.append(x)
    for y in graph[x]:
        if y == parent[x]:
            continue
        parent[y] = x
        stack.append(y)
        depth[y] = depth[x] + 1


def merge(v, p):
    A = dp[v]
    B = dp[p]
    size_A = size[v]
    size_B = size[p]
    c = C[p]
    if len(A) < len(B):
        A, B = B, A
        size_A, size_B = size_B, size_A
    answer[c] += (size_A - A[c]) * (size_B - B[c])
    # B を A にマージ
    for key, cnt in B.items():
        answer[key] -= cnt * A[key]
        A[key] += cnt
    size[p] = size_A + size_B
    dp[p] = A


answer = [0] * (N + 1)
dp = [defaultdict(int) for _ in range(N + 1)]
size = [0] * (N + 1)

for v in order[::-1]:
    c = C[v]
    n = size[v] + 1 - dp[v][c]
    answer[c] -= size[v] * n
    size[v] += 1
    dp[v][c] = size[v]
    if v == root:
        break
    p = parent[v]
    merge(v, p)

A = dp[root]
c = size[root]
for v in range(1, N + 1):
    print(c * A[v] + answer[v])

Submission status: Accepted

Solution code:
import sys
input = lambda: sys.stdin.readline().rstrip()
N = int(input())
C = [int(a) - 1 for a in input().split()]
X = [[] for i in range(N)]
for i in range(N-1):
    x, y = map(int, input().split())
    x, y = x-1, y-1
    X[x].append(y)
    X[y].append(x)

def EulerTour(n, X, i0):
    def f(k):
        return k * (k + 1) // 2
    USED = [0] * n
    USEDORG = [0] * n
    USEDTMP = [0] * n
    P = [-1] * n
    Q = [~i0, i0]
    ct = -1
    ET = []
    ET1 = [0] * n
    ET2 = [0] * n
    DE = [0] * n
    ANS = [0] * n
    de = -1
    while Q:
        i = Q.pop()
        if i < 0:
            ET2[~i] = ct
            de -= 1
            USED[C[~i]] += 1 + USEDTMP[~i]
            USEDTMP[~i] = 0
            if ~i:
                p = P[~i]
                pc = C[p]
                k = ET2[~i] - ET1[~i] + 1 - USED[pc] + USEDORG[~i]
                ANS[pc] += f(k)
                USEDTMP[p] += k
            continue
        if i >= 0:
            if i:
                USEDORG[i] = USED[C[P[i]]]
            ET.append(i)
            ct += 1
            if ET1[i] == 0: ET1[i] = ct
            de += 1
            DE[i] = de
        for a in X[i][::-1]:
            if a != P[i]:
                P[a] = i
                for k in range(len(X[a])):
                    if X[a][k] == i:
                        del X[a][k]
                        break
                Q.append(~a)
                Q.append(a)
    for i in range(n):
        ANS[i] = f(n) - f(n - USED[i]) - ANS[i]
    return ANS

ANS = EulerTour(N, X, 0)
print(*ANS, sep = "\n")

Submission status: Accepted

Solution code:
import sys

sys.setrecursionlimit(2 * 10 ** 5)

N = int(input())
colors = [int(x) for x in input().split()]
edges = [[int(x) for x in input().split()] for i in range(N - 1)]
graph = [[] for i in range(N + 1)]
for u, v in edges:
    graph[u].append(v)
    graph[v].append(u)
parent = {1: None}
children = {}


def rootTree(node):
    children[node] = []
    for nbr in graph[node]:
        if nbr not in parent:
            parent[nbr] = node
            children[node].append(nbr)
            rootTree(nbr)


rootTree(1)


# Count the paths that don't contain a particular color
# If we can find a max c-less component (in the sense that the nodes are either leaves or bordered by color c) of size s, then there are (s * s + 1 / 2) c-less paths in that component
colorlessCount = [0 for c in range(N + 1)]
# In our DFS we want to increment the size of each component except for the one matching current color all at once but that will TLE
# Instead implicitly increment them by tracking current size as dfs time `t` minus some offset. Then we can increment all just by adding to t or subtract from one by fiddling with the offset
colorlessTime = [0 for c in range(N + 1)]
t = 0


def buildComps(node):
    global t
    currColor = colors[node - 1]
    # Increment all by 1 except currColor
    t += 1
    colorlessTime[currColor] += 1
    for child in children[node]:
        # Each child forms a new currColor-less component so we need to save the old count
        oldCount = t - colorlessTime[currColor]
        # Start new comp
        colorlessTime[currColor] = t
        buildComps(child)
        # Done with child comp, count the paths and discard comp
        s = t - colorlessTime[currColor]
        colorlessCount[currColor] += s * (s + 1) // 2
        # Restore old count
        colorlessTime[currColor] = t - oldCount


buildComps(1)

# Count the colorless components path from root too
for c in range(1, N + 1):
    s = t - colorlessTime[c]
    colorlessCount[c] += s * (s + 1) // 2

# We counted colorless but want colored paths so take complement
maxPairs = N * (N + 1) // 2
for c in range(1, N + 1):
    print(maxPairs - colorlessCount[c])