Dynamic Segment Tree (ds/dynsegtree.hpp)

Description

Dynamic segment tree is a space-efficient segment tree that only allocates nodes when needed. It can handle range queries with point updates under a monoid in $\mathcal O(\log n)$ time per operation, using $\mathcal O(q \log n)$ space where $q$ is the number of operations.

Example monoid:

struct Sum {
  using T = int;
  static T id() { return 0; }
  static T op(T l, T r) { return l + r; }
};

Operations

DynSegTree<M, I>(int n)
- Construct dynamic segment tree with coordinate space [0, n)
- Time complexity: $\mathcal O(1)$
void set(I p, T x)
- Sets the value at p to be x
- Time complexity: $\mathcal O(\log n)$
T operator[](I p) or T get(I p)
- Returns the value at p, returns identity if not set
- Time complexity: $\mathcal O(\log n)$
T operator()(I l, I r) or T prod(I l, I r)
- Returns the product of the range [l, r)
- Time complexity: $\mathcal O(\log n)$
T all_prod()
- Returns the product of the entire coordinate space
- Time complexity: $\mathcal O(\log n)$

Depends on

other/st_alloc.hpp

Verified with

test/yosupo/ds/segment_tree/point_set_range_composite_large_array.test.cpp

Code

#pragma once

#include "../other/st_alloc.hpp"

// clang-format off
template<class M, class I = int> struct DynSegTree {
    using T = typename M::T;
    DynSegTree() = default;
    DynSegTree(int _n) : n(_n), root(new node()) {}

    void set(I p, const T &x) const {
        assert(0 <= p && p < n);
        auto rec = [&](auto &self, ptr t, I l, I r) -> void {
            if (r - l == 1) return void(t->val = x);
            I m = (l + r) / 2;
            if (p < m) self(self, t->c(0), l, m);
            else self(self, t->c(1), m, r);
            t->val = M::op(val(t->l), val(t->r));
        };
        rec(rec, root, 0, n);
    }
    T operator[](I p) const {
        assert(0 <= p && p < n);
        auto rec = [&](auto &self, ptr t, I l, I r) -> T {
            if (!t) return M::id();
            if (r - l == 1) return t->val;
            I m = (l + r) / 2;
            return p < m ? self(self, t->l, l, m) : self(self, t->r, m, r);
        };
        return rec(rec, root, 0, n);
    }
    T get(I p) const { return (*this)[p]; }
    T operator()(I l, I r) const {
        assert(0 <= l && l <= r && r <= n);
        auto rec = [&](auto &self, ptr t, I tl, I tr) -> T {
            if (!t) return M::id();
            if (l <= tl && tr <= r) return t->val;
            I tm = (tl + tr) / 2;
            if (r <= tm) return self(self, t->l, tl, tm);
            if (tm <= l) return self(self, t->r, tm, tr);
            return M::op(self(self, t->l, tl, tm), self(self, t->r, tm, tr));
        };
        return rec(rec, root, 0, n);
    }
    T prod(I l, I r) const { return (*this)(l, r); }
    T all_prod() const { return root->val; }

  private:
    struct node;
    using ptr = node *;
    struct node : st_alloc<node> {
        T val = M::id();
        ptr l = nullptr, r = nullptr;
        node() {}
        inline ptr c(bool z) {
            return !z ? l ? l : l = new node() : r ? r : r = new node();
        }
    };

    I n;
    ptr root;

    static inline T val(ptr t) { return t ? t->val : M::id(); }
};
// clang-format on

#line 2 "ds/dynsegtree.hpp"

#line 2 "other/st_alloc.hpp"

template<class T> struct st_alloc {
    static inline stack<T> pool;
    void *operator new(size_t) { return pool.emplace(), &pool.top(); }
    void operator delete(void *) {}
};
#line 4 "ds/dynsegtree.hpp"

// clang-format off
template<class M, class I = int> struct DynSegTree {
    using T = typename M::T;
    DynSegTree() = default;
    DynSegTree(int _n) : n(_n), root(new node()) {}

    void set(I p, const T &x) const {
        assert(0 <= p && p < n);
        auto rec = [&](auto &self, ptr t, I l, I r) -> void {
            if (r - l == 1) return void(t->val = x);
            I m = (l + r) / 2;
            if (p < m) self(self, t->c(0), l, m);
            else self(self, t->c(1), m, r);
            t->val = M::op(val(t->l), val(t->r));
        };
        rec(rec, root, 0, n);
    }
    T operator[](I p) const {
        assert(0 <= p && p < n);
        auto rec = [&](auto &self, ptr t, I l, I r) -> T {
            if (!t) return M::id();
            if (r - l == 1) return t->val;
            I m = (l + r) / 2;
            return p < m ? self(self, t->l, l, m) : self(self, t->r, m, r);
        };
        return rec(rec, root, 0, n);
    }
    T get(I p) const { return (*this)[p]; }
    T operator()(I l, I r) const {
        assert(0 <= l && l <= r && r <= n);
        auto rec = [&](auto &self, ptr t, I tl, I tr) -> T {
            if (!t) return M::id();
            if (l <= tl && tr <= r) return t->val;
            I tm = (tl + tr) / 2;
            if (r <= tm) return self(self, t->l, tl, tm);
            if (tm <= l) return self(self, t->r, tm, tr);
            return M::op(self(self, t->l, tl, tm), self(self, t->r, tm, tr));
        };
        return rec(rec, root, 0, n);
    }
    T prod(I l, I r) const { return (*this)(l, r); }
    T all_prod() const { return root->val; }

  private:
    struct node;
    using ptr = node *;
    struct node : st_alloc<node> {
        T val = M::id();
        ptr l = nullptr, r = nullptr;
        node() {}
        inline ptr c(bool z) {
            return !z ? l ? l : l = new node() : r ? r : r = new node();
        }
    };

    I n;
    ptr root;

    static inline T val(ptr t) { return t ? t->val : M::id(); }
};
// clang-format on