Lomsat gelral (Codeforces 600E)
Tags: Others, dsu

You are given a rooted tree with root in vertex 1. Each vertex is coloured in some colour.Let's call colour c dominating in the subtree of vertex v if there are no other colours that appear in the subtree of vertex v more times than colour c. So it's possible that two or more colours will be dominating in the subtree of some vertex.The subtree of vertex v is the vertex v and all other vertices that contains vertex v in each path to the root.For each vertex v find the sum of all dominating colours in the subtree of vertex v.

Input: The first line contains integer n (1 ≤ n ≤ 105) — the number of vertices in the tree.The second line contains n integers ci (1 ≤ ci ≤ n), ci — the colour of the i-th vertex.Each of the next n - 1 lines contains two integers xj, yj (1 ≤ xj, yj ≤ n) — the edge of the tree. The first vertex is the root of the tree.

Output: Print n integers — the sums of dominating colours for each vertex.


Sample Inputs:
4
1 2 3 4
1 2
2 3
2 4

15
1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
1 2
1 3
1 4
1 14
1 15
2 5
2 6
2 7
3 8
3 9
3 10
4 11
4 12
4 13

Sample Outputs:
10 9 3 4

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

Solution (Python):
import sys
sys.setrecursionlimit(2 * 10 ** 5)
input = sys.stdin.readline
from collections import defaultdict, deque, Counter
from heapq import heappop, heappush
from bisect import bisect_left, bisect_right
from math import gcd
import threading
threading.stack_size(2 ** 26)

def main():
    n = int(input())
    c = list(map(int, input().split()))
    graph = [[] for i in range(n + 1)]
    for i in range(n - 1):
        (a, b) = map(int, input().split())
        graph[a].append(b)
        graph[b].append(a)
    ans = [0] * n

    def dfs(crr, prev):
        d = defaultdict(int)
        dcnt = dict()
        col = c[crr - 1]
        d[col] += 1
        dcnt[1] = col
        for nxt in graph[crr]:
            if nxt == prev:
                continue
            (nxtd, nxtdcnt) = dfs(nxt, crr)
            if len(nxtd) <= len(d):
                for (k, v) in nxtd.items():
                    pre = d[k]
                    d[k] += v
                    if pre in dcnt:
                        dcnt[pre] -= k
                        if dcnt[pre] == 0:
                            dcnt.pop(pre)
                    if d[k] not in dcnt:
                        dcnt[d[k]] = k
                    else:
                        dcnt[d[k]] += k
            else:
                for (k, v) in d.items():
                    pre = nxtd[k]
                    nxtd[k] += v
                    if pre in nxtdcnt:
                        nxtdcnt[pre] -= k
                        if nxtdcnt[pre] == 0:
                            nxtdcnt.pop(pre)
                    if nxtd[k] not in nxtdcnt:
                        nxtdcnt[nxtd[k]] = k
                    else:
                        nxtdcnt[nxtd[k]] += k
                d = nxtd
                dcnt = nxtdcnt
        mx = [-1, -1]
        for (k, v) in dcnt.items():
            mx = max([k, v], mx)
        ans[crr - 1] = mx[1]
        return (d, dcnt)
    dfs(1, -1)
    print(*ans)
threading.Thread(target=main).start()

Solution (Python):
import sys, math
import threading
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 8)

def main():

    def f(a, b, c, s):
        if len(a) < len(b):
            t = b
            b = a
            a = t
        for x in b:
            if x in a:
                a[x] += b[x]
            else:
                a[x] = b[x]
            if a[x] > c:
                c = a[x]
                s = 0
            if a[x] == c:
                s += x
        return [a, c, s]

    def dfs(v, x, ans, a, gp):
        if v[x] == 1:
            return [{}, 0, 0]
        d = {a[x - 1]: 1}
        v[x] = 1
        c = 1
        s = a[x - 1]
        for i in gp[x]:
            t = dfs(v, i, ans, a, gp)
            if len(t[0]) > len(d):
                s = t[2]
                c = t[1]
            t = f(d, t[0], c, s)
            (d, c, s) = (t[0], t[1], t[2])
        ans[x] = s
        return [d, c, s]
    for _ in range(1):
        n = int(input())
        a = list(map(int, input().split()))
        gp = []
        for i in range(n + 1):
            gp.append([])
        for i in range(n - 1):
            (x, y) = map(int, input().split())
            gp[x].append(y)
            gp[y].append(x)
        v = [0] * (n + 1)
        ans = [0] * (n + 1)
        dfs(v, 1, ans, a, gp)
        ans = ans[1:]
        print(*ans)
threading.stack_size(10 ** 8)
t = threading.Thread(target=main)
t.start()
t.join()

Solution (Python):
import threading
import sys
from sys import stdin
input = stdin.readline
sys.setrecursionlimit(10 ** 8)
from collections import defaultdict

def main():
    n = int(input())
    arr = list(map(int, input().split()))
    arr.insert(0, 0)
    dit = defaultdict(list)
    for i in range(n - 1):
        (a, b) = map(int, input().split())
        dit[a].append(b)
        dit[b].append(a)
    ans = [0] * n

    def fun(node, par):
        mp = [defaultdict(int)]
        i = 0
        (fq, sm) = (0, 0)
        for e in dit[node]:
            if e != par:
                (x, f, s) = fun(e, node)
                if len(x) > len(mp[i]):
                    i = len(mp)
                    fq = f
                    sm = s
                mp.append(x)
        for k in range(len(mp)):
            if k == i:
                continue
            for e in mp[k]:
                mp[i][e] += mp[k][e]
                if mp[i][e] > fq:
                    fq = mp[i][e]
                    sm = e
                elif mp[i][e] == fq:
                    sm += e
        e = arr[node]
        mp[i][e] += 1
        if mp[i][e] > fq:
            fq = mp[i][e]
            sm = e
        elif mp[i][e] == fq:
            sm += e
        ans[node - 1] = sm
        return (mp[i], fq, sm)
    fun(1, -1)
    print(*ans)
threading.stack_size(10 ** 8)
t = threading.Thread(target=main)
t.start()
t.join()