disjoint_set/segment_tree.hpp

#include <functional>
#include <vector>

template <typename T, typename U>
class SegmentTree
{
    int n;
    std::function<void(T &, U &)> _update;
    std::function<T(T &, T &)> _merge;
    std::function<void(T &, T &, T &, int, int, int)> _push;

    void update(int p, int b, int e, int x, int y, U &upd)
    {
        if (x <= b && e <= y)
        {
            _update(ds[p], upd);
        }
        else
        {
            int m = (b + e) >> 1, l = p << 1, r = l | 1;

            _push(ds[p], ds[l], ds[r], b, m, e);

            if (x < m)
                update(l, b, m, x, y, upd);

            if (m < y)
                update(r, m, e, x, y, upd);

            ds[p] = _merge(ds[l], ds[r]);
        }
    }

    T query(int p, int b, int e, int x, int y)
    {
        if (x <= b && e <= y)
        {
            return ds[p];
        }
        else
        {
            int m = (b + e) >> 1, l = p << 1, r = l | 1;

            _push(ds[p], ds[l], ds[r], b, m, e);

            if (x < m)
            {
                auto le = query(l, b, m, x, y);

                if (m < y)
                {
                    auto ri = query(r, m, e, x, y);
                    return _merge(le, ri);
                }
                else
                {
                    return le;
                }
            }
            else
            {
                return query(r, m, e, x, y);
            }
        }
    }

public:
    std::vector<T> ds;

    SegmentTree(int n,
                std::function<void(T &, U &)> update,
                std::function<T(T &, T &)> merge,
                std::function<void(T &, T &, T &, int, int, int)> push)
        : n(n), _update(update), _merge(merge), _push(push)
    {
        ds = std::vector<T>(4 * n);
    }

    // Update on point
    void update(int x, U upd)
    {
        update(1, 0, n, x, x + 1, upd);
    }

    // Update on range
    void update(int x, int y, U upd)
    {
        update(1, 0, n, x, y, upd);
    }

    // Query on a point
    T query(int x)
    {
        return query(1, 0, n, x, x + 1);
    }

    // Query on a range
    T query(int x, int y)
    {
        return query(1, 0, n, x, y);
    }
};

template <typename T>
void noop_push(T &p, T &l, T &r, int b, int m, int e) {}

template <typename T, typename U>
class SegmentTreeBuilder
{
    std::function<void(T &, U &)> _update;
    std::function<T(T &, T &)> _merge;
    std::function<void(T &, T &, T &, int, int, int)> _push;

public:
    SegmentTreeBuilder() : _push(noop_push<T>)
    {
    }

    SegmentTree<T, U> *init(int size)
    {
        return new SegmentTree<T, U>(size, _update, _merge, _push);
    }

    SegmentTreeBuilder<T, U> *with_lazy(std::function<void(T &, T &, T &, int, int, int)> push)
    {
        _push = push;
        return this;
    }

    SegmentTreeBuilder<T, U> *with_lazy(std::function<void(T &, T &, T &)> push)
    {

        _push = [push](T &p, T &l, T &r, int b, int m, int e) {
            push(p, l, r);
        };
        return this;
    }

    SegmentTreeBuilder<T, U> *with_update(std::function<void(T &, U &)> update)
    {
        _update = update;
        return this;
    }

    SegmentTreeBuilder<T, U> *with_merge(std::function<T(T &, T &)> merge)
    {
        _merge = merge;
        return this;
    }
};

template <typename T, typename U, typename R = T>
SegmentTreeBuilder<T, U> *segment_tree()
{
    return new SegmentTreeBuilder<T, U>();
}

// TODO: Allow custom initializer
// TODO: Use array of length 2 * n (instead of 4 * n)
// TODO: Use iterative segment tree whenever possible.