Li Hua and Tree (Codeforces 1797D)
Tags: Implement, trees

Li Hua has a tree of $$$n$$$ vertices and $$$n-1$$$ edges. The root of the tree is vertex $$$1$$$. Each vertex $$$i$$$ has importance $$$a_i$$$. Denote the size of a subtree as the number of vertices in it, and the importance as the sum of the importance of vertices in it. Denote the heavy son of a non-leaf vertex as the son with the largest subtree size. If multiple of them exist, the heavy son is the one with the minimum index.Li Hua wants to perform $$$m$$$ operations:  "1 $$$x$$$" ($$$1\leq x \leq n$$$) — calculate the importance of the subtree whose root is $$$x$$$.  "2 $$$x$$$" ($$$2\leq x \leq n$$$) — rotate the heavy son of $$$x$$$ up. Formally, denote $$$son_x$$$ as the heavy son of $$$x$$$, $$$fa_x$$$ as the father of $$$x$$$. He wants to remove the edge between $$$x$$$ and $$$fa_x$$$ and connect an edge between $$$son_x$$$ and $$$fa_x$$$. It is guaranteed that $$$x$$$ is not root, but not guaranteed that $$$x$$$ is not a leaf. If $$$x$$$ is a leaf, please ignore the operation. Suppose you were Li Hua, please solve this problem.

Input: The first line contains 2 integers $$$n,m$$$ ($$$2\le n\le 10^{5},1\le m\le 10^{5}$$$) — the number of vertices in the tree and the number of operations.The second line contains $$$n$$$ integers $$$a_{1},a_{2},\ldots ,a_{n}$$$ ($$$-10^{9}\le a_{i}\le 10^{9}$$$) — the importance of each vertex.Next $$$n-1$$$ lines contain the edges of the tree. The $$$i$$$-th line contains two integers $$$u_i$$$ and $$$v_i$$$ ($$$1\le u_i,v_i\le n$$$, $$$u_i\ne v_i$$$) — the corresponding edge. The given edges form a tree.Next $$$m$$$ lines contain operations — one operation per line. The $$$j$$$-th operation contains two integers $$$t_{j},x_{j}$$$ ($$$t_{j}\in \{1,2\}$$$, $$$1 \leq x_{j} \leq n$$$, $$$x_{j}\neq 1$$$ if $$$t_j = 2$$$) — the $$$j$$$-th operation. 

Output: For each query "1 $$$x$$$", output the answer in an independent line.

Note: In the first example:The initial tree is shown in the following picture:  The importance of the subtree of $$$6$$$ is $$$a_6+a_7=2$$$.After rotating the heavy son of $$$3$$$ (which is $$$6$$$) up, the tree is shown in the following picture:  The importance of the subtree of $$$6$$$ is $$$a_6+a_3+a_7=3$$$.The importance of the subtree of $$$2$$$ is $$$a_2+a_4+a_5=3$$$.

Sample Inputs:
7 4
1 1 1 1 1 1 1
1 2
1 3
2 4
2 5
3 6
6 7
1 6
2 3
1 6
1 2

10 14
-160016413 -90133231 -671446275 -314847579 -910548234 121155052 -359359950 83112406 -704889624 145489303
1 6
1 10
10 8
1 4
3 4
2 7
2 5
3 2
9 8
1 4
2 2
2 4
1 4
1 10
2 10
1 9
1 6
2 8
2 10
1 5
1 8
1 1
2 5

Sample Outputs:
2
3
3

-2346335269
-314847579
-476287915
-704889624
121155052
-1360041415
228601709
-2861484545

Solution (Python):
import collections
import functools
import heapq
import itertools
import sys
RD = sys.stdin.readline
INT = lambda : int(RD())
STRS = lambda : RD().split()
INTS = lambda : list(map(int, RD().split()))
T = 1
for _ in range(T):
    (n, m) = INTS()
    a = INTS()
    e = [set() for i in range(n)]
    for _ in range(n - 1):
        (u, v) = INTS()
        u -= 1
        v -= 1
        e[u].add(v)
        e[v].add(u)
    s = [0] * n
    h = [-1] * n
    f = [-1] * n
    sz = [0] * n

    def dfs(u, fa=-1):
        f[u] = fa
        s[u] += a[u]
        sz[u] = 1
        for v in e[u]:
            if v == fa:
                continue
            dfs(v, u)
            s[u] += s[v]
            sz[u] += sz[v]

    def bfs(u, fa=-1):
        q = [(u, fa)]
        i = 0
        while i < len(q):
            (x, fx) = q[i]
            i += 1
            f[x] = fx
            s[x] = a[x]
            sz[x] = 1
            for u in e[x]:
                if u != fx:
                    q.append((u, x))
        for i in range(len(q) - 1, 0, -1):
            (x, fx) = q[i]
            s[fx] += s[x]
            sz[fx] += sz[x]
    bfs(0, -1)
    hs = [[] for i in range(n)]

    def buildhs():
        for u in range(n):
            for v in e[u]:
                if v != f[u]:
                    hs[u].append((-sz[v], v))
            heapq.heapify(hs[u])
    buildhs()

    def isleaf(x):
        return len(e[x]) == 1

    def geth(x):
        if h[x] != -1:
            return h[x]
        q = hs[x]
        assert q
        while q:
            (usz, u) = q[0]
            if u not in e[x] or u == f[x] or -usz != sz[u]:
                heapq.heappop(q)
                continue
            h[x] = u
            break
        return h[x]

    def rotate(x):
        if isleaf(x):
            return
        fx = f[x]
        hx = geth(x)
        s[x] -= s[hx]
        s[hx] += s[x]
        sz[x] -= sz[hx]
        sz[hx] += sz[x]
        e[fx].remove(x)
        e[x].remove(fx)
        e[hx].add(fx)
        e[fx].add(hx)
        f[x] = hx
        f[hx] = fx
        h[fx] = h[x] = h[hx] = -1
        heapq.heappush(hs[hx], (-sz[x], x))
        heapq.heappush(hs[fx], (-sz[hx], hx))
    for _ in range(m):
        (t, x) = INTS()
        x -= 1
        if t == 1:
            print(s[x])
        else:
            rotate(x)

Solution (Python):
import collections
import functools
import heapq
import itertools
import sys
RD = sys.stdin.readline
INT = lambda : int(RD())
STRS = lambda : RD().split()
INTS = lambda : list(map(int, RD().split()))
T = 1
for _ in range(T):
    (n, m) = INTS()
    a = INTS()
    e = [set() for i in range(n)]
    for _ in range(n - 1):
        (u, v) = INTS()
        u -= 1
        v -= 1
        e[u].add(v)
        e[v].add(u)
    s = [0] * n
    h = [-1] * n
    f = [-1] * n
    sz = [0] * n

    def dfs(u, fa=-1):
        f[u] = fa
        s[u] += a[u]
        sz[u] = 1
        for v in e[u]:
            if v == fa:
                continue
            dfs(v, u)
            s[u] += s[v]
            sz[u] += sz[v]

    def bfs(u, fa=-1):
        q = [(u, fa)]
        i = 0
        while i < len(q):
            (x, fx) = q[i]
            i += 1
            f[x] = fx
            s[x] = a[x]
            sz[x] = 1
            for u in e[x]:
                if u != fx:
                    q.append((u, x))
        for i in range(len(q) - 1, 0, -1):
            (x, fx) = q[i]
            s[fx] += s[x]
            sz[fx] += sz[x]
    bfs(0, -1)
    hs = [[] for i in range(n)]

    def buildhs():
        for u in range(n):
            for v in e[u]:
                if v != f[u]:
                    hs[u].append((-sz[v], v))
            heapq.heapify(hs[u])
    buildhs()

    def isleaf(x):
        return len(e[x]) == 1

    def geth(x):
        if h[x] != -1:
            return h[x]
        q = hs[x]
        assert q
        while q:
            (usz, u) = q[0]
            if u not in e[x] or u == f[x] or -usz != sz[u]:
                heapq.heappop(q)
                continue
            h[x] = u
            break
        return h[x]

    def rotate(x):
        if isleaf(x):
            return
        fx = f[x]
        hx = geth(x)
        s[x] -= s[hx]
        s[hx] += s[x]
        sz[x] -= sz[hx]
        sz[hx] += sz[x]
        e[fx].remove(x)
        e[x].remove(fx)
        e[hx].add(fx)
        e[fx].add(hx)
        f[x] = hx
        f[hx] = fx
        h[fx] = h[x] = h[hx] = -1
        heapq.heappush(hs[hx], (-sz[x], x))
        heapq.heappush(hs[fx], (-sz[hx], hx))
    for _ in range(m):
        (t, x) = INTS()
        x -= 1
        if t == 1:
            print(s[x])
        else:
            rotate(x)

Solution (Python):
import sys
import math
from heapq import heappush, heappop
from collections import defaultdict
input = sys.stdin.readline

def inp():
    return int(input())

def inlt():
    return list(map(int, input().split()))

def insr():
    s = input()
    return list(s[:len(s) - 1])

def invr():
    return list(map(int, input().split()))
if __name__ == '__main__':
    (n, m) = inlt()
    a = [0] + inlt()
    g = defaultdict(list)
    for _ in range(n - 1):
        (u, v) = inlt()
        g[u].append(v)
        g[v].append(u)
    t = [1]
    p = {1: 0}
    i = 0
    while i < len(t):
        u = t[i]
        for v in g[u]:
            if v in p:
                continue
            t.append(v)
            p[v] = u
        i += 1
    sz = [1] * (n + 1)
    ss = [[] for i in range(n + 1)]
    for v in t[:0:-1]:
        u = p[v]
        sz[u] += sz[v]
        a[u] += a[v]
        heappush(ss[u], (-sz[v], v))
    for _ in range(m):
        (ty, x) = inlt()
        if ty == 1:
            print(a[x])
        elif sz[x] > 1:
            u = p[x]
            (mx, v) = heappop(ss[x])
            while p[v] != x or mx != -sz[v]:
                (mx, v) = heappop(ss[x])
            p[v] = u
            p[x] = v
            a[x] -= a[v]
            sz[x] -= sz[v]
            a[v] += a[x]
            sz[v] += sz[x]
            heappush(ss[u], (-sz[v], v))
            heappush(ss[v], (-sz[x], x))
    pass