Editorial for Pondo Subtree Add

Use this editorial only when stuck, and do not 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$ and record an enter/exit Euler tour: store $\mathrm{first}[v]$ when the search enters $v$ , and $\mathrm{last}[v]$ when it leaves. The subtree of $v$ is then the contiguous range $[\mathrm{first}[v], \mathrm{last}[v]]$ on the tour, and the value of vertex $u$ lives at index $\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$ to $[L, R]$ is the two point updates $d[L] \gets d[L] + x$ and $d[R + 1] \gets d[R + 1] - x$ , and the value at index $i$ is the prefix sum $d[0] + \cdots + d[i]$ . Initialise each $a_u$ as a range add of $a_u$ on $[\mathrm{first}[u], \mathrm{first}[u]]$ .

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

Solution (Python)

Code 1

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

Comments0

No comments yet

Be the first to comment.

Editorial for Pondo Subtree Add

Approach

Solution (Python)

Comments0

New comment