Buggy python submission (verdict: Runtime Error):
import asyncio

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


async def go(curr, fro):
    col = Cs[curr] - 1
    deg = 1
    clust = {}
    for n in adj[curr]:
        if n == fro:
            continue
        ch_deg, ch_clust = await 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


_, root_clust = asyncio.get_event_loop().run_until_complete(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)


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

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


stack = [[None, None], [go(0, -1), None]]
while len(stack) > 1:
    gen, val = stack[-1]
    try:
        params = gen.send(val)
    except StopIteration as e:
        stack[-2][1] = e.value
        stack.pop()
    else:
        stack.append([go(*params), None])

_, root_clust = stack[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)


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