Editorial for Pondo Subtree Add

Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

Approach

Root the tree at vertex $1$ ~1~ and record an enter/exit Euler tour: store $\mathrm{first}[v]$ ~\mathrm{first}[v]~ when the search enters $v$ ~v~, and $\mathrm{last}[v]$ ~\mathrm{last}[v]~ when it leaves. The subtree of $v$ ~v~ is then the contiguous range $[\mathrm{first}[v], \mathrm{last}[v]]$ ~[\mathrm{first}[v], \mathrm{last}[v]]~ on the tour, and the value of vertex $u$ ~u~ lives at index $\mathrm{first}[u]$ ~\mathrm{first}[u]~.

A subtree add becomes a range add on that interval, and a vertex query becomes a point query. Maintain a difference array in a Fenwick tree: adding $x$ ~x~ to $[L, R]$ ~[L, R]~ is the two point updates $d[L] \gets d[L] + x$ ~d[L] \gets d[L] + x~ and $d[R + 1] \gets d[R + 1] - x$ ~d[R + 1] \gets d[R + 1] - x~, and the value at index $i$ ~i~ is the prefix sum $d[0] + \cdots + d[i]$ ~d[0] + \cdots + d[i]~. Initialise each $a_u$ ~a_u~ as a range add of $a_u$ ~a_u~ on $[\mathrm{first}[u], \mathrm{first}[u]]$ .

The tour and Fenwick operations together run in $O((n + q) \log n)$ ~O((n + q) \log n)~.

Solution (Python)

import sys

sys.setrecursionlimit(1_000_000)


class Fenwick:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)

    def add(self, i, delta):
        while i <= self.n:
            self.bit[i] += delta
            i += i & -i

    def prefix(self, i):
        s = 0
        while i > 0:
            s += self.bit[i]
            i -= i & -i
        return s

    def range_add(self, left, right, delta):
        self.add(left + 1, delta)
        self.add(right + 2, -delta)

    def point(self, i):
        return self.prefix(i + 1)


data = list(map(int, sys.stdin.buffer.read().split()))
it = iter(data)
n = next(it)
q = next(it)
a = [next(it) for _ in range(n)]
adj = [[] for _ in range(n)]
for _ in range(n - 1):
    u = next(it) - 1
    v = next(it) - 1
    adj[u].append(v)
    adj[v].append(u)

first = [0] * n
last = [0] * n
timer = 0


def dfs(u, p):
    global timer
    first[u] = timer
    timer += 1
    for v in adj[u]:
        if v != p:
            dfs(v, u)
    last[u] = timer
    timer += 1


dfs(0, -1)
fw = Fenwick(timer + 2)
for u in range(n):
    fw.range_add(first[u], first[u], a[u])

out = []
for _ in range(q):
    t = next(it)
    if t == 1:
        v = next(it) - 1
        x = next(it)
        fw.range_add(first[v], last[v], x)
    else:
        u = next(it) - 1
        out.append(str(fw.point(first[u])))

if out:
    print("\n".join(out))

Comments

There are no comments at the moment.