Alex's Anthology of Algorithms Common Code for Contests in Concise C++
Data Structures / Range Queries in Two Dimensions

2.6.2 Sparse 2D Segment Tree

2-Data-Structures/2.6.2_Sparse_2D_Segment_Tree.cpp

Maintain a two-dimensional array over a huge grid while supporting dynamic queries of rectangular subarrays and dynamic updates of individual indices. This is a sparse (a.k.a. dynamic or implicit) 2D segment tree: row and column nodes are allocated lazily as cells are touched, so large coordinate bounds are supported without allocating the full grid. Inner trees path-compress chains of unallocated column nodes, so each point update creates only $O(1)$ column nodes at every visited row node. Allocated nodes are kept in stable-address pools and released together when the tree is destroyed.

The query operation is defined by a commutative associative aggregate function combine(a, b). Because untouched regions are implicit, combine_n(v, area) must return the aggregate summary of area copies of the initial value v. The default code below assumes a numerical array type, defining queries for the sum of the target range. For rectangle-min queries, combine(a, b) should return std::min(a, b) and combine_n(v, area) should return v.

The point update operation is defined by apply_delta(v, d), which returns the new value at one updated cell. The default code below defines increments. For point assignment, apply_delta(v, d) should return d.

Rows and columns are split independently, so every rectangle query decomposes into $O(\log\left(R\right) \cdot \log\left(C\right))$ canonical rectangles regardless of whether it is thin, off-center, or otherwise adversarially placed.

Use the dense 2D segment tree in the previous section when the grid is modest enough for $O(R \cdot C)$ storage; it is simpler and faster. Use this sparse version when the coordinate range is huge but only a small fraction of cells are updated. For dense additive rectangle sums, prefer the 2D Fenwick tree in 2.7.5.

  • SparseSegTree2D<T, R, C>(v = T{}) constructs a two-dimensional array over rows $[0, {\htmlClass{math-inline-code}{\texttt{R}}})$ and columns $[0, {\htmlClass{math-inline-code}{\texttt{C}}})$. All array values are implicitly initialized to v. Nodes are allocated lazily as indices are touched.
  • at(r, c) returns the value at row r, column c.
  • query(r1, c1, r2, c2) returns the result of combine() applied to every value in the rectangular region consisting of rows in $[{\htmlClass{math-inline-code}{\texttt{r1}}}, {\htmlClass{math-inline-code}{\texttt{r2}}}]$ and columns in $[{\htmlClass{math-inline-code}{\texttt{c1}}}, {\htmlClass{math-inline-code}{\texttt{c2}}}]$.
  • update(r, c, d) adds d to the value at (r, c).

Overflow warning: Products of initial values with rectangle areas, and all resulting sums, must fit in T.

Time Complexity
  • $O(1)$ per call to the constructor.
  • $O(\log\left(R\right) \cdot \log\left(C\right))$ per call to at(), query(), and update().
Space Complexity
  • $O(n \log R)$ for storage after $n$ point updates.
  • $O(\log\left(R\right) + \log\left(C\right))$ auxiliary stack space for at(), query(), and update().

Implementation

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <deque>
#include <optional>

template<typename T, int R = 1000000001, int C = 1000000001>
class SparseSegTree2D {
  static_assert(R > 0 && C > 0);

  static T combine(const T &a, const T &b) { return a + b; }
  static T combine_n(const T &v, int64_t area) { return v * area; }
  static T apply_delta(const T &v, const T &d) { return v + d; }

  struct InnerNode {
    T value;
    int lo, hi;
    InnerNode *left, *right;

    InnerNode(int lo, int hi, const T &v)
        : value(v), lo(lo), hi(hi), left(nullptr), right(nullptr) {}
  };

  struct OuterNode {
    InnerNode inner;
    int lo, hi;
    OuterNode *left, *right;

    OuterNode(int lo, int hi, const T &v)
        : inner(0, C - 1, v), lo(lo), hi(hi), left(nullptr), right(nullptr) {}
  };

  std::deque<InnerNode> inner_nodes;
  std::deque<OuterNode> outer_nodes;
  OuterNode *root;
  T init;

  static int64_t length(int lo, int hi) { return hi - lo + 1LL; }
  static void append_result(std::optional<T> &res, const T &v) { res = res ? combine(*res, v) : v; }

  InnerNode *make_inner(int lo, int hi, const T &v) {
    inner_nodes.emplace_back(lo, hi, v);
    return &inner_nodes.back();
  }

  OuterNode *make_outer(int lo, int hi, const T &v) {
    outer_nodes.emplace_back(lo, hi, v);
    return &outer_nodes.back();
  }

  template<typename Node, typename Get>
  T query_nodes(const Node *n, int qlo, int qhi, int64_t span, const Get &get) const {
    if (n == nullptr) {
      return combine_n(init, span * length(qlo, qhi));
    }
    int lo = n->lo, hi = n->hi, mid = lo + (hi - lo) / 2;
    std::optional<T> res;
    if (qlo < lo) {
      append_result(res, combine_n(init, span * length(qlo, std::min(qhi, lo - 1))));
    }
    int ql = std::max(qlo, lo), qr = std::min(qhi, hi);
    if (ql <= qr) {
      if (ql == lo && qr == hi) {
        append_result(res, get(n));
      } else {
        if (ql <= mid) {
          append_result(res, query_nodes(n->left, ql, std::min(qr, mid), span, get));
        }
        if (mid < qr) {
          append_result(res, query_nodes(n->right, std::max(ql, mid + 1), qr, span, get));
        }
      }
    }
    if (hi < qhi) {
      append_result(res, combine_n(init, span * length(std::max(qlo, hi + 1), qhi)));
    }
    return *res;
  }

  T query_inner(const InnerNode *n, int c1, int c2, int64_t rows) const {
    return query_nodes(n, c1, c2, rows, [](const InnerNode *node) { return node->value; });
  }

  T query_outer(const OuterNode *n, int r1, int r2, int c1, int c2) const {
    return query_nodes(n, r1, r2, length(c1, c2), [&](const OuterNode *node) {
      return query_inner(&node->inner, c1, c2, length(node->lo, node->hi));
    });
  }

  template<typename Apply>
  void update_inner(InnerNode *n, int c, int64_t rows, const Apply &apply) {
    int lo = n->lo, hi = n->hi, mid = lo + (hi - lo) / 2;
    if (lo == hi) {
      n->value = apply(n->value);
      return;
    }
    InnerNode *&target = (c <= mid) ? n->left : n->right;
    if (target == nullptr) {
      target = make_inner(c, c, combine_n(init, rows));
    }
    if (target->lo <= c && c <= target->hi) {
      update_inner(target, c, rows, apply);
    } else {
      int split_lo = lo, split_hi = hi, split_mid = mid;
      do {
        if (c <= split_mid) {
          split_hi = split_mid;
        } else {
          split_lo = split_mid + 1;
        }
        split_mid = split_lo + (split_hi - split_lo) / 2;
      } while ((c <= split_mid) == (target->lo <= split_mid));
      InnerNode *tmp =
          make_inner(split_lo, split_hi, combine_n(init, rows * length(split_lo, split_hi)));
      (target->lo <= split_mid ? tmp->left : tmp->right) = target;
      target = tmp;
      update_inner(tmp, c, rows, apply);
    }
    n->value =
        combine(query_inner(n->left, lo, mid, rows), query_inner(n->right, mid + 1, hi, rows));
  }

  void update(OuterNode *n, int r, int c, const T &d) {
    int lo = n->lo, hi = n->hi, mid = lo + (hi - lo) / 2;
    int64_t rows = length(lo, hi);
    if (lo == hi) {
      update_inner(&n->inner, c, 1, [&](const T &v) { return apply_delta(v, d); });
      return;
    }
    if (r <= mid) {
      if (n->left == nullptr) {
        n->left = make_outer(lo, mid, combine_n(init, length(lo, mid) * C));
      }
      update(n->left, r, c, d);
    } else {
      if (n->right == nullptr) {
        n->right = make_outer(mid + 1, hi, combine_n(init, length(mid + 1, hi) * C));
      }
      update(n->right, r, c, d);
    }
    T value = combine_n(init, rows);
    if (n->left != nullptr || n->right != nullptr) {
      T lval = (n->left != nullptr) ? query_inner(&n->left->inner, c, c, length(lo, mid))
                                    : combine_n(init, length(lo, mid));
      T rval = (n->right != nullptr) ? query_inner(&n->right->inner, c, c, length(mid + 1, hi))
                                     : combine_n(init, length(mid + 1, hi));
      value = combine(lval, rval);
    }
    update_inner(&n->inner, c, rows, [&](const T &) { return value; });
  }

 public:
  explicit SparseSegTree2D(const T &v = T{}) : root(nullptr), init(v) {
    root = make_outer(0, R - 1, combine_n(v, static_cast<int64_t>(R) * C));
  }

  SparseSegTree2D(const SparseSegTree2D &) = delete;
  SparseSegTree2D &operator=(const SparseSegTree2D &) = delete;
  T at(int r, int c) const { return query(r, c, r, c); }

  T query(int r1, int c1, int r2, int c2) const {
    assert(0 <= r1 && r1 <= r2 && r2 < R);
    assert(0 <= c1 && c1 <= c2 && c2 < C);
    return query_outer(root, r1, r2, c1, c2);
  }

  void update(int r, int c, const T &d) {
    assert(0 <= r && r < R && 0 <= c && c < C);
    update(root, r, c, d);
  }
};

Example Usage

#include <iostream>
using namespace std;

int main() {
  SparseSegTree2D<int> t(0);
  t.update(0, 0, 7);
  t.update(0, 1, 6);
  t.update(1, 0, 5);
  t.update(1, 1, 4);
  t.update(2, 1, 1);
  t.update(2, 2, 9);
  cout << "Values:" << endl;
  for (int i = 0; i < 3; i++) {
    for (int j = 0; j < 3; j++) {
      cout << t.at(i, j) << " ";
    }
    cout << endl;
  }
  assert(t.query(0, 0, 0, 1) == 13);
  assert(t.query(0, 0, 1, 0) == 12);
  assert(t.query(1, 1, 2, 2) == 14);
  assert(t.query(0, 0, 1000000000, 1000000000) == 32);
  t.update(500000000, 500000000, -100);
  assert(t.query(0, 0, 1000000000, 1000000000) == -68);

  SparseSegTree2D<int, 1, 4> rectangular(0);
  rectangular.update(0, 0, 5);
  rectangular.update(0, 3, 6);
  assert(rectangular.query(0, 0, 0, 3) == 11);

  SparseSegTree2D<int, 2, 3> initialized(2);
  assert(initialized.query(0, 0, 1, 2) == 12);
  initialized.update(0, 1, 3);
  assert(initialized.query(0, 0, 1, 2) == 15);
  return 0;
}

Example Output

Values:
7 6 0
5 4 0
0 1 9
Back to top