Editorial for Range Updates


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

Use a difference array. If d is the difference array of x, then:

  • x_k = d_1 + d_2 + \cdots + d_k
  • adding u to the range [l, r] becomes two point updates: d_l \gets d_l + u and d_{r + 1} \gets d_{r + 1} - u

Store d in a segment tree (or Fenwick tree) that supports point updates and prefix-sum queries. Each range update is then two point updates, and each point query is one prefix sum.

The total time complexity is O((n + q) \log n).

Solution (Python)

from __future__ import annotations

from typing import Callable, Generic, List, TypeVar
import copy

T = TypeVar("T")


class SegmentTree(Generic[T]):
    def __init__(self, arr: List[T], combine: Callable[[T, T], T]) -> None:
        self.combine = combine
        self.length = len(arr)
        self.array = [0] * (self.length * 4 + 1)
        self.__build(1, 0, self.length - 1, arr)

    def __build(self, v: int, tl: int, tr: int, arr: List[T]) -> None:
        if tl == tr:
            self.array[v] = arr[tl]
        else:
            tm = (tl + tr) // 2
            self.__build(2 * v, tl, tm, arr)
            self.__build(2 * v + 1, tm + 1, tr, arr)
            self.array[v] = self.combine(self.array[2 * v], self.array[2 * v + 1])

    def update(self, index: int, value: T) -> None:
        self.__update(index, 1, 0, self.length - 1, value)

    def __update(self, index: int, v: int, tl: int, tr: int, value: T) -> None:
        if tl == tr:
            self.array[v] = value
            return

        tm = (tl + tr) // 2
        if index <= tm:
            self.__update(index, 2 * v, tl, tm, value)
        else:
            self.__update(index, 2 * v + 1, tm + 1, tr, value)

        self.array[v] = self.combine(self.array[2 * v], self.array[2 * v + 1])

    def query(self, l: int, r: int) -> T:
        return copy.copy(self.__query(1, 0, self.length - 1, l, r))

    def __query(self, v: int, tl: int, tr: int, l: int, r: int) -> T:
        if tl == l and tr == r:
            return self.array[v]

        tm = (tl + tr) // 2
        if r <= tm:
            return self.__query(2 * v, tl, tm, l, r)
        if l > tm:
            return self.__query(2 * v + 1, tm + 1, tr, l, r)
        return self.combine(
            self.__query(2 * v, tl, tm, l, tm),
            self.__query(2 * v + 1, tm + 1, tr, tm + 1, r),
        )


n = int(input())
arr = [int(x) for x in input().split()]

arr2 = [0] * (n + 1)
for i in range(n):
    arr2[i] += arr[i]
    arr2[i + 1] -= arr[i]


def combine(a: int, b: int) -> int:
    return a + b


tree = SegmentTree[int](arr2, combine)

for _ in range(int(input())):
    query = input()
    if query[0] == "1":
        _, a, b, u = (int(x) for x in query.split())
        tree.update(a - 1, tree.query(a - 1, a - 1) + u)
        tree.update(b, tree.query(b, b) - u)
    else:
        _, k = (int(x) for x in query.split())
        print(tree.query(0, k - 1))

Comments

There are no comments at the moment.