Jamie and Tree (Codeforces 916E)
Tags: Others, trees

To your surprise, Jamie is the final boss! Ehehehe.Jamie has given you a tree with n vertices, numbered from 1 to n. Initially, the root of the tree is the vertex with number 1. Also, each vertex has a value on it.Jamie also gives you three types of queries on the tree:1 v — Change the tree's root to vertex with number v.2 u v x — For each vertex in the subtree of smallest size that contains u and v, add x to its value.3 v — Find sum of values of vertices in the subtree of vertex with number v.A subtree of vertex v is a set of vertices such that v lies on shortest path from this vertex to root of the tree. Pay attention that subtree of a vertex can change after changing the tree's root.Show your strength in programming to Jamie by performing the queries accurately!

Input: The first line of input contains two space-separated integers n and q (1 ≤ n ≤ 105, 1 ≤ q ≤ 105) — the number of vertices in the tree and the number of queries to process respectively.The second line contains n space-separated integers a1, a2, ..., an ( - 108 ≤ ai ≤ 108) — initial values of the vertices.Next n - 1 lines contains two space-separated integers ui, vi (1 ≤ ui, vi ≤ n) describing edge between vertices ui and vi in the tree.The following q lines describe the queries.Each query has one of following formats depending on its type:1 v (1 ≤ v ≤ n) for queries of the first type.2 u v x (1 ≤ u, v ≤ n,  - 108 ≤ x ≤ 108) for queries of the second type.3 v (1 ≤ v ≤ n) for queries of the third type.All numbers in queries' descriptions are integers.The queries must be carried out in the given order. It is guaranteed that the tree is valid.

Output: For each query of the third type, output the required answer. It is guaranteed that at least one query of the third type is given by Jamie.

Note: The following picture shows how the tree varies after the queries in the first sample.   

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

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

Sample Outputs:
27
19
5

18
21

Solution (Python):
import sys
from types import GeneratorType
input = lambda : sys.stdin.readline().strip()
print = lambda d: sys.stdout.write(str(d) + '\n')
M = lambda : map(int, input().split())
read = lambda : list(M())
N = 10 ** 5 + 10
T = 17
fa = [[0] * (T + 1) for _ in range(N)]
child = [[] for _ in range(N)]
deep = [0] * N
(s, e, dfn) = (deep.copy(), deep.copy(), deep.copy())
cnt = 0
(bit1, bit2) = (deep.copy(), deep.copy())
lowBit = [0] * N
LOG2 = [0] * N

def bootstrap(f, stack=[]):

    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

def add(i: int, v: int):
    x = i
    while i <= n:
        bit1[i] += v
        bit2[i] += (x - 1) * v
        i += lowBit[i]

def query(i: int):
    ans = 0
    x = i
    while i:
        ans += x * bit1[i] - bit2[i]
        i -= lowBit[i]
    return ans

def Update(l: int, r: int, x: int):
    add(l, x)
    add(r + 1, -x)

def Query(l: int, r: int):
    return query(r) - query(l - 1)
(n, q) = M()
val = [0] + read()
for i in range(1, n + 1):
    lowBit[i] = i & -i
for i in range(2, n + 1):
    LOG2[i] = LOG2[i >> 1] + 1
for _ in range(n - 1):
    (u, v) = M()
    child[u].append(v)
    child[v].append(u)

@bootstrap
def dfs(curr: int, parent: int):
    fa[curr][0] = parent
    deep[curr] = deep[parent] + 1
    global cnt
    cnt += 1
    dfn[cnt] = curr
    s[curr] = cnt
    for i in range(1, T + 1):
        fa[curr][i] = fa[fa[curr][i - 1]][i - 1]
    for nxt in child[curr]:
        if nxt != parent:
            yield dfs(nxt, curr)
    e[curr] = cnt
    yield

def lca(x: int, y: int):
    if deep[x] < deep[y]:
        (x, y) = (y, x)
    for i in range(T, -1, -1):
        if deep[fa[x][i]] >= deep[y]:
            x = fa[x][i]
    if x == y:
        return x
    for i in range(T, -1, -1):
        if fa[x][i] != fa[y][i]:
            x = fa[x][i]
            y = fa[y][i]
    return fa[x][0]
root = 1

def LCA(x, y):
    (t1, t2, t3) = (lca(root, x), lca(root, y), lca(x, y))
    mxDeep = max(deep[t1], deep[t2], deep[t3])
    if mxDeep == deep[t1]:
        return t1
    if mxDeep == deep[t2]:
        return t2
    return t3

def top(x: int, k: int):
    while k:
        step = LOG2[k]
        x = fa[x][step]
        k -= 1 << step
    return x

def Add(curr: int, v: int):
    if curr == root:
        Update(1, n, v)
    elif s[curr] <= s[root] and e[root] <= e[curr]:
        t = top(root, deep[root] - deep[curr] - 1)
        Update(1, n, v)
        Update(s[t], e[t], -v)
    else:
        Update(s[curr], e[curr], v)

def Ans(curr: int):
    if curr == root:
        return Query(1, n)
    elif s[curr] <= s[root] and e[root] <= e[curr]:
        t = top(root, deep[root] - deep[curr] - 1)
        return Query(1, n) - Query(s[t], e[t])
    return Query(s[curr], e[curr])
dfs(1, 0)
for i in range(1, n + 1):
    t = val[dfn[i]] - val[dfn[i - 1]]
    bit1[i] += t
    bit2[i] += (i - 1) * t
    j = i + lowBit[i]
    if j <= n:
        bit1[j] += bit1[i]
        bit2[j] += bit2[i]
for _ in range(q):
    op = read()
    if op[0] == 1:
        root = op[1]
    elif op[0] == 2:
        (u, v, val) = op[1:]
        Add(LCA(u, v), val)
    else:
        print(Ans(op[1]))

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 1)
(dfs, hashing, read_from_file) = (0, 0, 0)
de = 1
if 1:
    if standard_input:
        import io, os, sys
        input = lambda : sys.stdin.readline().strip()
        inf = float('inf')

        def I():
            return input()

        def II():
            return int(input())

        def MII():
            return map(int, input().split())

        def LI():
            return list(input().split())

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

        def LFI():
            return list(map(float, input().split()))

        def GMI():
            return map(lambda x: int(x) - 1, input().split())

        def LGMI():
            return list(map(lambda x: int(x) - 1, input().split()))
    if packages:
        from io import BytesIO, IOBase
        import math
        import random
        import os
        import bisect
        import typing
        from collections import Counter, defaultdict, deque
        from copy import deepcopy
        from functools import cmp_to_key, lru_cache, reduce
        from heapq import merge, heapify, heappop, heappush, heappushpop, nlargest, nsmallest
        from itertools import accumulate, combinations, permutations, count, product
        from operator import add, iand, ior, itemgetter, mul, xor
        from string import ascii_lowercase, ascii_uppercase, ascii_letters
        from typing import *
        BUFSIZE = 4096
    if output_together:

        class FastIO(IOBase):
            newlines = 0

            def __init__(self, file):
                self._fd = file.fileno()
                self.buffer = BytesIO()
                self.writable = 'x' in file.mode or 'r' not in file.mode
                self.write = self.buffer.write if self.writable else None

            def read(self):
                while True:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    if not b:
                        break
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines = 0
                return self.buffer.read()

            def readline(self):
                while self.newlines == 0:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    self.newlines = b.count(b'\n') + (not b)
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines -= 1
                return self.buffer.readline()

            def flush(self):
                if self.writable:
                    os.write(self._fd, self.buffer.getvalue())
                    (self.buffer.truncate(0), self.buffer.seek(0))

        class IOWrapper(IOBase):

            def __init__(self, file):
                self.buffer = FastIO(file)
                self.flush = self.buffer.flush
                self.writable = self.buffer.writable
                self.write = lambda s: self.buffer.write(s.encode('ascii'))
                self.read = lambda : self.buffer.read().decode('ascii')
                self.readline = lambda : self.buffer.readline().decode('ascii')
        sys.stdout = IOWrapper(sys.stdout)
    if dfs:
        from types import GeneratorType

        def bootstrap(f, stack=[]):

            def wrappedfunc(*args, **kwargs):
                if stack:
                    return f(*args, **kwargs)
                else:
                    to = f(*args, **kwargs)
                    while True:
                        if type(to) is GeneratorType:
                            stack.append(to)
                            to = next(to)
                        else:
                            stack.pop()
                            if not stack:
                                break
                            to = stack[-1].send(to)
                    return to
            return wrappedfunc
    if hashing:
        RANDOM = random.getrandbits(20)

        class Wrapper(int):

            def __init__(self, x):
                int.__init__(x)

            def __hash__(self):
                return super(Wrapper, self).__hash__() ^ RANDOM
    if read_from_file:
        file = open('input.txt', 'r').readline().strip()[1:-1]
        fin = open(file, 'r')
        input = lambda : fin.readline().strip()
        output_file = open('output.txt', 'w')

        def fprint(*args, **kwargs):
            print(*args, **kwargs, file=output_file)
    if de:

        def debug(*args, **kwargs):
            print('\x1b[92m', end='')
            print(*args, **kwargs)
            print('\x1b[0m', end='')

class LCA:

    def __init__(self, depth, parent):
        self.n = len(depth)
        self.root = 0
        self.num = self.n.bit_length()
        self.depth = depth
        self.parent = [[-1] * self.n for i in range(self.num)]
        self.parent[0] = parent
        for k in range(self.num - 1):
            for v in range(self.n):
                if self.parent[k][v] == -1:
                    self.parent[k + 1][v] = -1
                else:
                    self.parent[k + 1][v] = self.parent[k][self.parent[k][v]]

    def getLCA(self, u, v):
        if self.depth[u] > self.depth[v]:
            (u, v) = (v, u)
        for k in range(self.num):
            if self.depth[v] - self.depth[u] >> k & 1:
                v = self.parent[k][v]
        if u == v:
            return u
        for k in reversed(range(self.num)):
            if self.parent[k][u] != self.parent[k][v]:
                u = self.parent[k][u]
                v = self.parent[k][v]
        return self.parent[0][u]

    def search(self, v, x):
        for k in reversed(range(self.num)):
            if x >> k & 1:
                v = self.parent[k][v]
        return v

class FenwickTree:

    def __init__(self, nums):
        n = len(nums)
        self.n = n
        self.bit = nums
        for i in range(n):
            if i | i + 1 < n:
                self.bit[i | i + 1] += self.bit[i]

    def sum(self, r):
        res = 0
        while r >= 0:
            res += self.bit[r]
            r = (r & r + 1) - 1
        return res

    def add(self, idx, delta):
        while idx < self.n:
            self.bit[idx] += delta
            idx = idx | idx + 1

def main():
    (n, q) = MII()
    nums = LII()
    path = [[] for _ in range(n)]
    for _ in range(n - 1):
        (u, v) = GMI()
        path[u].append(v)
        path[v].append(u)
    left = [0] * n
    right = [0] * n
    parent = [-1] * n
    depth = [0] * n
    dfs_order = []
    stack = [0]
    tmstamp = 0
    while stack:
        u = stack.pop()
        if u >= 0:
            dfs_order.append(u)
            stack.append(~u)
            left[u] = tmstamp
            tmstamp += 1
            for v in path[u]:
                if parent[u] != v:
                    parent[v] = u
                    depth[v] = depth[u] + 1
                    stack.append(v)
        else:
            right[~u] = tmstamp
    vals = [nums[i] for i in dfs_order]
    k_tree = FenwickTree([0] * (n + 1))
    b_tree = FenwickTree([0] + vals)
    tot_sum = sum(vals)
    cur_add = 0
    lca = LCA(depth, parent)

    def add_range(l, r, x):
        nonlocal tot_sum
        tot_sum += (r - l) * x
        k_tree.add(l, x)
        k_tree.add(r, -x)
        b_tree.add(l, -l * x)
        b_tree.add(r, r * x)

    def sum_range(l, r):
        return k_tree.sum(r) * r + b_tree.sum(r) - k_tree.sum(l) * l - b_tree.sum(l) + (r - l) * cur_add
    root = 0
    for _ in range(q):
        query = LGMI()
        if query[0] == 0:
            root = query[1]
        elif query[0] == 1:
            (u, v, x) = (query[1], query[2], query[3] + 1)
            cnt = 0
            if left[root] <= left[u] and right[u] <= right[root]:
                cnt += 1
            if left[root] <= left[v] and right[v] <= right[root]:
                cnt += 1
            if cnt == 0:
                uv = lca.getLCA(u, v)
                if left[uv] <= left[root] and right[root] <= right[uv]:
                    pu = lca.getLCA(u, root)
                    pv = lca.getLCA(v, root)
                    chosen = pu if depth[pu] > depth[pv] else pv
                    cur_add += x
                    tot_sum += n * x
                    p = lca.search(root, depth[root] - depth[chosen] - 1)
                    add_range(left[p], right[p], -x)
                else:
                    add_range(left[uv], right[uv], x)
            elif cnt == 1:
                cur_add += x
                tot_sum += n * x
            else:
                l = lca.getLCA(u, v)
                if l == root:
                    cur_add += x
                    tot_sum += n * x
                else:
                    add_range(left[l], right[l], x)
        else:
            v = query[1]
            if v == root:
                print(tot_sum)
            elif left[v] <= left[root] and right[root] <= right[v]:
                v = lca.search(root, depth[root] - depth[v] - 1)
                print(tot_sum - sum_range(left[v], right[v]))
            else:
                print(sum_range(left[v], right[v]))
    return
t = 1
for _ in range(t):
    main()

Solution (Python):
import sys
input = lambda : sys.stdin.readline().rstrip()
sint = lambda : int(input())
mint = lambda : map(int, input().split())
ints = lambda : list(map(int, input().split()))

def solve() -> None:
    (n, m) = mint()
    nums = ints()
    g = [[] for _ in range(n)]
    for _ in range(n - 1):
        (u, v) = mint()
        g[u - 1].append(v - 1)
        g[v - 1].append(u - 1)
    k = n.bit_length()
    tin = [-1] * n
    tout = [-1] * n
    parent = [[-1] * n for _ in range(k)]
    depth = [-1] * n
    depth[0] = t = 0
    q = [0]
    order = []
    while q:
        u = q[-1]
        if tin[u] == -1:
            order.append(u)
            t += 1
            tin[u] = t
            for v in g[u]:
                if depth[v] >= 0:
                    continue
                depth[v] = depth[u] + 1
                parent[0][v] = u
                q.append(v)
        else:
            q.pop()
            tout[u] = t
    for u in order[::-1]:
        if parent[0][u] != -1:
            nums[parent[0][u]] += nums[u]
    for i in range(1, k):
        for j in range(n):
            if (p := parent[i - 1][j]) != -1:
                parent[i][j] = parent[i - 1][p]

    def go_up(u, k):
        for i in range(k.bit_length()):
            if k >> i & 1:
                u = parent[i][u]
        return u

    def lca(u, v):
        d = depth[u] - depth[v]
        if d >= 0:
            u = go_up(u, d)
        else:
            v = go_up(v, -d)
        if u == v:
            return u
        for p in range(k - 1, -1, -1):
            if parent[p][u] != parent[p][v]:
                (u, v) = (parent[p][u], parent[p][v])
        return parent[0][u]
    bit1 = [0] * (n + 1)
    bit2 = [0] * (n + 1)

    def query(x: int) -> int:
        (res, t) = (0, x)
        while x:
            res += bit1[x] * t - bit2[x]
            x &= x - 1
        return res

    def add(x: int, val: int):
        t = x
        while x <= n:
            bit1[x] += val
            bit2[x] += val * (t - 1)
            x += x & -x
    root = 0
    for _ in range(m):
        qry = ints()
        if qry[0] == 1:
            root = qry[1] - 1
        elif qry[0] == 2:
            (u, v, x) = (qry[1] - 1, qry[2] - 1, qry[3])
            anc = lca(u, v)
            if anc == root:
                add(1, x)
            elif not tin[anc] < tin[root] <= tout[anc]:
                add(tin[anc], x)
                add(tout[anc] + 1, -x)
            else:
                (lu, lv) = (lca(u, root), lca(v, root))
                mxd = max(depth[lu], depth[lv])
                if mxd < depth[root]:
                    w = go_up(root, depth[root] - mxd - 1)
                    add(tin[w], -x)
                    add(tout[w] + 1, x)
                add(1, x)
        else:
            u = qry[1] - 1
            if u == root:
                print(nums[0] + query(n))
            elif not tin[u] < tin[root] <= tout[u]:
                print(nums[u] + query(tout[u]) - query(tin[u] - 1))
            else:
                w = go_up(root, depth[root] - depth[u] - 1)
                print(nums[0] + query(n) - nums[w] - query(tout[w]) + query(tin[w] - 1))
solve()