Segment Tree

Implementations

C++

#include <concepts>
#include <cstddef>
#include <vector>
 
template <typename Ops, typename S, typename L>
concept LazySegConcept = requires(S a, S b, L f, L g, std::size_t len) {
	{ Ops::identity_value() } -> std::same_as<S>;
	{ Ops::combine(a, b) } -> std::same_as<S>;
	{ Ops::identity_lazy() } -> std::same_as<L>;
	{ Ops::compose(f, g) } -> std::same_as<L>;
	{ Ops::apply(f, a, len) } -> std::same_as<S>;
	{
		f == g
	}; // not really needed, just a minor checkf for where I check if no lazy
};
 
template <typename S, typename L, typename Ops>
    requires LazySegConcept<Ops, S, L>
class SegTree {
  private:
	using sz = std::size_t;
	sz n_ = 0;
	std::vector<S> tree_;
	std::vector<L> lazy_;
 
	inline sz left(sz node) { return node << 1; }
 
	inline sz right(sz node) { return (node << 1) + 1; }
 
	// build : O(n)
	void build(sz node, sz l, sz r, const std::vector<S> &values) {
		// [l,r) like range
		if (r - l == 1) {
			tree_[node] = values[l];
			return;
		}
 
		sz m = (l + r) >> 1;
		build(left(node), l, m, values);
		build(right(node), m, r, values);
		pull(node);
	}
 
	// helpers - pull, apply_lazy, push
	void pull(sz node) {
		tree_[node] = Ops::combine(tree_[left(node)], tree_[right(node)]);
	}
 
	void apply_lazy(L lazy, sz node, sz len) {
		// can swap these 2 lines
		// apply lazy s.t. it's green again
		tree_[node] = Ops::apply(lazy, tree_[node], len);
		lazy_[node] = Ops::compose(lazy, lazy_[node]);
	}
 
	void push(sz node, sz len) {
		if (lazy_[node] == Ops::identity_lazy() || len == 1)
			return;
 
		sz half = len >> 1;
		// make children green
		apply_lazy(lazy_[node], left(node), half);
		apply_lazy(lazy_[node], right(node), half);
		lazy_[node] = Ops::identity_lazy();
	}
 
	void set(sz node, sz l, sz r, sz i, const S &value) {
		if (r - l == 1) {
			tree_[node] = value;
			lazy_[node] = Ops::identity_lazy();
			return;
		}
 
		push(node, r - l);
		sz m = (l + r) >> 1;
		if (i < m)
			set(left(node), l, m, i, value);
		else
			set(right(node), m, r, i, value);
		pull(node);
	}
 
	S get(sz node, sz l, sz r, sz i) {
		if (r - l == 1) {
			return tree_[node];
		}
 
		push(node, r - l);
		sz m = (r + l) >> 1;
		if (i < m)
			return get(left(node), l, m, i);
		return get(right(node), m, r, i);
	}
 
	S query(sz node, sz l, sz r, sz ql, sz qr) {
		if (qr <= l || r <= ql)
			return Ops::identity_value();
		if (ql <= l && r <= qr)
			return tree_[node]; // what I enter into is ALWAYS green
 
		push(node, r - l);
		sz m = (l + r) >> 1;
		return Ops::combine(query(left(node), l, m, ql, qr),
		                    query(right(node), m, r, ql, qr));
	}
 
	void apply(sz node, sz l, sz r, sz ql, sz qr, const L &lazy) {
		if (qr <= l || r <= ql)
			return;
		if (ql <= l && r <= qr) {
			apply_lazy(lazy, node, r - l);
			return;
		}
 
		push(node, r - l); // make children green
		sz m = (l + r) >> 1;
		apply(left(node), l, m, ql, qr, lazy);
		apply(right(node), m, r, ql, qr, lazy);
		pull(node);
	}
 
  public:
	// a type alias is not a distinct type
	// this'll work with std:size_t args
	explicit SegTree(sz n)
	    : n_(n), tree_(4 * n + 1, Ops::identity_value()),
	      lazy_(4 * n + 1, Ops::identity_lazy()) {}
 
	explicit SegTree(const std::vector<S> &values) : SegTree(values.size()) {
		if (n_ > 0)
			build(1, 0, n_, values);
	}
 
	void set(sz i, const S &value) { set(1, 0, n_, i, value); }
 
	S get(sz i) { return get(1, 0, n_, i); }
 
	S query(sz l, sz r) { return query(1, 0, n_, l, r); }
 
	void apply(sz l, sz r, const L &lazy) { apply(1, 0, n_, l, r, lazy); }
};
 
struct RangeAddSum {
	static long long identity_value() { return 0; }
	static long long combine(long long a, long long b) { return a + b; }
	static long long identity_lazy() { return 0; }
	static long long compose(long long newer, long long older) {
		return newer + older;
	}
	static long long apply(long long lazy, long long value, std::size_t len) {
		return value + lazy * static_cast<long long>(len);
	}
};
 
int main() {
	SegTree<long long, long long, RangeAddSum> st({1, 2, 3, 4});
	st.apply(1, 3, 5);
	return static_cast<int>(st.query(0, 4));
}

Java

import java.util.*;
 
interface LazySegOps<S, L> {
    S identityValue();
 
    L identityLazy();
 
    S combine(S a, S b);
 
    L compose(L f, L g);
 
    S apply(L f, S a, int len);
}
 
final class SegTree<S, L> {
    private final LazySegOps<S, L> ops;
    private int n;
    private ArrayList<S> tree;
    private ArrayList<L> lazy;
 
    public SegTree(LazySegOps<S, L> ops, int n) {
        this.ops = ops;
        // NOTE: remember the 4n+1
        tree = new ArrayList<>(Collections.nCopies(4 * n + 1, ops.identityValue()));
        lazy = new ArrayList<>(Collections.nCopies(4 * n + 1, ops.identityLazy()));
        this.n = n;
    }
 
    private void build(int node, int l, int r, List<S> list) {
        if (r - l == 1) {
            tree.set(node, list.get(l));
            return;
        }
 
        int m = (l + r) >> 1;
        build(left(node), l, m, list);
        build(right(node), m, r, list);
        pull(node);
    }
 
    public void build(List<S> list) {
        build(1, 0, n, list);
    }
 
    // helpers
    private int left(int node) {
        return node << 1;
    }
 
    private int right(int node) {
        return (node << 1) + 1;
    }
 
    private void pull(int node) {
        tree.set(node, ops.combine(tree.get(left(node)), tree.get(right(node))));
    }
 
    private void applyLazy(L lazyVal, int node, int len) {
        tree.set(node, ops.apply(lazyVal, tree.get(node), len));
        lazy.set(node, ops.compose(lazyVal, lazy.get(node)));
    }
 
    private void push(int node, int len) {
        // NOTE: remember to use the .equals here and not ==
        // also note that this assumes lazy cannot be null since then this equal may
        // throw
        if (len == 1 || lazy.get(node).equals(ops.identityLazy()))
            return;
 
        int half = len >> 1;
        applyLazy(lazy.get(node), left(node), half);
        applyLazy(lazy.get(node), right(node), half);
        lazy.set(node, ops.identityLazy());
    }
 
    // actual funcs
    private S get(int node, int l, int r, int i) {
        if (r - l == 1) {
            return tree.get(node);
        }
 
        int m = (l + r) >> 1;
        push(node, r - l);
        if (i < m)
            return get(left(node), l, m, i);
        return get(right(node), m, r, i);
    }
 
    private void set(int node, int l, int r, int i, S val) {
        if (r - l == 1) {
            // it is an actual hardcoded SET
            tree.set(node, val);
            lazy.set(node, ops.identityLazy());
            return;
        }
 
        push(node, r - l);
        int m = (l + r) >> 1;
        if (i < m)
            set(left(node), l, m, i, val);
        else
            set(right(node), m, r, i, val);
        pull(node);
    }
 
    private S query(int node, int l, int r, int ql, int qr) {
        // REM: THIS IS WHAT'S NATURAL TO ME
        if (qr <= l || r <= ql)
            return ops.identityValue();
        if (ql <= l && r <= qr)
            return tree.get(node);
 
        push(node, r - l);
        int m = (l + r) >> 1;
        return ops.combine(query(left(node), l, m, ql, qr), query(right(node), m, r, ql, qr));
    }
 
    private void apply(int node, int l, int r, int ql, int qr, L lazyVal) {
        if (qr <= l || r <= ql)
            return;
        if (ql <= l && r <= qr) {
            applyLazy(lazyVal, node, r - l);
            return;
        }
 
        push(node, r - l);
        int m = (l + r) >> 1;
        apply(left(node), l, m, ql, qr, lazyVal);
        apply(right(node), m, r, ql, qr, lazyVal);
        pull(node);
    }
 
    public S get(int i) {
        return get(1, 0, n, i);
    }
 
    public void set(int i, S val) {
        set(1, 0, n, i, val);
    }
 
    public S query(int l, int r) {
        return query(1, 0, n, l, r);
    }
 
    public void apply(int l, int r, L lazyVal) {
        apply(1, 0, n, l, r, lazyVal);
    }
}

Notes

core

this is in handwritten notes ( I need to digitise them )

cpp