Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter
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)

dp = [Counter() 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

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        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')


Fixed submission:
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')

Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter

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

adj = [list() 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)


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)

dp = [Counter() 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

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        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')


Fixed submission:
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')

Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter

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

adj = [list() 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)

ROOT = N//2

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

degree = [0]*N
stack = [ROOT]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    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

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

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

total = (subtreesize[ROOT]+1)*(subtreesize[ROOT]+2)//2

print(*(total-v for v in cnt), sep='\n')


Fixed submission:
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')

Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter

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

adj = [list() 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)

ROOT = N//2

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

degree = [0]*N
stack = [ROOT]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    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

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

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

total = (subtreesize[ROOT]+1)*(subtreesize[ROOT]+2)//2

print(*(total-v for v in cnt), sep='\n')


Fixed submission:
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')

Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter

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

adj = [list() 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)


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

degree = [0]*N
stack = [0]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    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

        a,b = dp[p],dp[c]
        if len(a) > len(b):
            a += b
            dp[p] = a
        else:
            b += a
            dp[p] = b
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

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

total = (subtreesize[0]+1)*(subtreesize[0]+2)//2

print(*(total-v for v in cnt), sep='\n')


Fixed submission:
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')

Buggy python submission (verdict: Time Limit Exceeded):
from collections import Counter

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

adj = [list() 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)


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

degree = [0]*N
stack = [0]
while stack:
    p = stack.pop()
    for c in adj[p]:
        if parent[c] is None:
            parent[c] = p
            degree[p] += 1
            stack.append(c)

leaves = [i for i in range(N) if degree[i] == 0]

dp = [Counter() for _ in range(N)] # color-wise counter
subtreesize = [0]*N
cnt = [0]*N

while leaves:
    c = leaves.pop()
    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] += dp[c]
        dp[c] = None
        degree[p] -= 1
        if degree[p] == 0:
            leaves.append(p)

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

total = (subtreesize[0]+1)*(subtreesize[0]+2)//2

print(*(total-v for v in cnt), sep='\n')


Fixed submission:
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')