Segment-Tree (Readme: https://shininglionking.blogspot.tw/2018/02/segment-tree.html)

SegmentTree.cpp

/*
 * =====================================================================================
 *
 *       Filename:  SegmentTree.cpp
 *
 *    Description:  Implementation of the data structure: segment tree
 *
 *        Version:  1.0
 *        Created:  2018/02/18 (yyyy/mm/dd)
 *       Revision:  none
 *       Compiler:  g++14
 *
 *         Author:  lionking
 *   Organization:  None
 *
 * =====================================================================================
 */

#include <cstdio>
#include <vector>
#include <iostream>
#define CHOICE 0 // enter a number to choose an example


// test whether the given type has overloaded operator() or not
// concept: SFINAE
template <class T>
struct has_function
{
    typedef char true_type;
    typedef long false_type;

    // template not available if there's no overloaded operator< in U's scope
    // Note: operator() must be non-static member function
    template <class U>
    static true_type test(typeof (&U::operator()));

    // fallback
    template <class U>
    static false_type test(...);

    // tests which overload of test is chosen for T
    static const bool value = sizeof(test<T>(0)) == sizeof(true_type);
};


// pre-declaration
template <bool b>
struct DefaultSegTreeOpImpl;


// specialization: situation when the given type exist overloaded operator()
template <>
struct DefaultSegTreeOpImpl<true>
{
    template <typename T>
    T operator() (const T& lhs, const T& rhs) const
    {
        T op;
        return op(lhs, rhs);
    }
};


// specialization: situation when the given type do not have overloaded operator()
template <>
struct DefaultSegTreeOpImpl<false>
{
    template <typename T>
    T operator() (const T& lhs, const T& rhs) const { return std::min(lhs, rhs); }
};


struct DefaultSegTreeOp
{
    template <typename T>
    T operator() (const T& lhs, const T& rhs) const
    {
        DefaultSegTreeOpImpl<has_function<T>::value> op;
        return op(lhs, rhs);
    }
};


/*
 * Tree node for constructing segment tree
 */
template <typename T>
struct SegTreeNode
{
    size_t first, last; // specify this tree node can cover the original data within [first, last)
    T data;

    SegTreeNode(size_t l=0, size_t r=0) : first(l), last(r) {};
    SegTreeNode(const T& new_data, size_t l=0, size_t r=0) : first(l), last(r), data(new_data) {};
    void set(const T& new_data, size_t l, size_t r) { data = new_data; first = l; last = r; }
};


/********************************************************************************************************
 * Segment tree
 * Given an array of data, this class will construct a segment tree for query in O(log |N|),
 * where N is the number of data in the given array.
 *
 * template parameter:
 * 1. T: type of the given data
 * 2. Operator is a binary function that accepts two arguments with type T and returns a result in type T
 *    default: std::min
 *
 * Note:
 *   When T is a user-defined structure, it must overload at least one of the following operators:
 *   1. operator<
 *   2. operator()
 *   When both of the above operators are overloaded, operator() will be chosen for execution.
 ********************************************************************************************************/
template < typename T, typename Operator=DefaultSegTreeOp >
class SegmentTree
{
public:
    typedef T value_type;
    typedef Operator op_type;


    // parameters:
    // first, last:
    // Random-access iterators to the initial and final positions of the sequence to be constructed in segment tree.
    // The range used is [first,last), which contains all the elements between first and last,
    // including the element pointed by first but not the element pointed by last.
    template <typename RandomAccessIterator>
    void construct(RandomAccessIterator first, RandomAccessIterator last);


    // query data within range [first, last)
    // Note: first and last are indices in original data set,
    // and index starts from 0
    inline T query(const size_t first, const size_t last) const;


    // update data on the specified index
    // Note: parameter index is the index in original data set and starts from 0
    inline void update(const size_t index, const T& data);


private:
    std::vector< SegTreeNode<T> > tree;
    Operator compare;
    static constexpr size_t tree_root = 0;


    inline size_t getlc(const size_t node) const { return ((node << 1) + 1); }
    inline size_t getrc(const size_t node) const { return ((node + 1) << 1); }


    // sub-function for construct (segment tree construction)
    template <typename RandomAccessIterator>
    void construct(const RandomAccessIterator& first, size_t curr, size_t start, size_t end);


    // query data within range [first, last)
    // curr: currntly queried tree node
    T query(const size_t curr, const size_t first, const size_t last) const;


    // update segment tree
    // curr: currntly updated tree node
    void update(const size_t index, const T& data, const size_t curr);
};


template <typename T, typename Operator>
template <typename RandomAccessIterator>
void SegmentTree<T, Operator>::construct(const RandomAccessIterator& first, size_t curr, size_t start, size_t end)
{
    if (start == (end - 1)) {
        tree[curr].set(*(first + start), start, end);
    }
    else if (start < end) {
        const auto middle = (start + end) >> 1;
        const auto lchild = getlc(curr);
        const auto rchild = getrc(curr);

        construct(first, lchild, start, middle); // recursive to left child
        construct(first, rchild, middle, end); // recursive to right child
        tree[curr].set(compare(tree[lchild].data, tree[rchild].data), start, end);
    }
    else { return; }
}


template <typename T, typename Operator>
template <typename RandomAccessIterator>
void SegmentTree<T, Operator>::construct(RandomAccessIterator first, RandomAccessIterator last)
{
    const size_t size = std::distance(first, last);

    // Minimum required leaf nodes is the minimum x s.t. 2^x >= number of elements
    size_t tree_size = 1;
    for (; (tree_size < size); tree_size <<= 1);
    // Then, minimum required tree node is 2^(x+1)
    tree.resize(tree_size << 1);

    // segment tree construction
    construct(first, tree_root, 0, size);
}


template <typename T, typename Operator>
T SegmentTree<T, Operator>::query(const size_t curr, const size_t first, const size_t last) const
{
    const auto &curr_node = tree.at(curr);
    const auto lchild = getlc(curr);
    const auto rchild = getrc(curr);

    if ((first >= curr_node.last) || (last < curr_node.first)) {
        // current node cannot cover query range.
        return T();
    }
    else if ((first <= curr_node.first) && (curr_node.last <= last)) {
        // query range has fully cover current node
        // return data stored in current node;
        return curr_node.data;
    }
    else if (last <= tree.at(lchild).last) {
        // left child can fully cover query range
        return query(lchild, first, last);
    }
    else if (tree.at(rchild).first <= first) {
        // right chile can fully cover query range
        return query(rchild, first, last);
    }
    else {
        const auto& lhs = query(lchild, first, std::min(last, tree.at(lchild).last));
        const auto& rhs = query(rchild, std::max(first, tree.at(rchild).first), last);
        return compare(lhs, rhs);
    }
}


template <typename T, typename Operator>
inline T SegmentTree<T, Operator>::query(const size_t first, const size_t last) const
{
    return query(tree_root, first, last);
}


template <typename T, typename Operator>
void SegmentTree<T, Operator>::update(const size_t index, const T& data, const size_t curr)
{
    auto& curr_node = tree.at(curr);

    if ((index < curr_node.first) || (curr_node.last <= index)) {
        // tree nodes that should be updated is not within the range corresponding to currently processed node
        return;
    }
    else if ((index == curr_node.first) && ((curr_node.last - curr_node.first) == 1)) {
        // target leaf node, update data directly
        curr_node.data = data;
    }
    else {
        const auto lchild = getlc(curr);
        const auto rchild = getrc(curr);
        update(index, data, lchild);    // update data in left child
        update(index, data, rchild);    // update data in right child
        // merge: update data in current node
        curr_node.data = compare(tree.at(lchild).data, tree.at(rchild).data);
    }
}


template <typename T, typename Operator>
inline void SegmentTree<T, Operator>::update(const size_t index, const T& data)
{
    update(index, data, tree_root);
}


/****************************************************
 * The followings are examples of using SegmentTree *
 ****************************************************/


struct SegTreeOp1
{
    template <typename T>
    T operator() (const T& lhs, const T& rhs) const { return std::max(lhs, rhs); }
};


struct SegTreeOp2
{
    template <typename T>
    T operator() (const T& lhs, const T& rhs) const { return (lhs + rhs); }
};


struct STNode1
{
    int value;

    STNode1(int v = 0) : value(v) {}

    bool operator< (const STNode1& rhs) const { return value < rhs.value; }
};


struct STNode2
{
    int value;

    STNode2(int v = 0) : value(v) {}

    STNode2 operator() (const STNode2& lhs, const STNode2& rhs) const { return std::max(lhs.value, rhs.value); }
};


struct STNode3
{
    int value;

    STNode3(int v = 0) : value(v) {}

    bool operator< (const STNode3& rhs) const { return value < rhs.value; }
    STNode3 operator() (const STNode3& lhs, const STNode3& rhs) const { return std::max(lhs.value, rhs.value); }
};


struct STNode4
{
    int value;

    STNode4(int v = 0) : value(v) {}
};


bool operator< (const STNode4& lhs, const STNode4& rhs) { return lhs.value < rhs.value; }


int main()
{
    std::cout << "choice: " << CHOICE << '\n';
#if CHOICE == 0
    // The simplest example:
    // 1. Use primitive type
    // 2. No extra comparator is specified for segment tree
    const int data[] = {1, 5, 2, 3, 4};
    SegmentTree<int> st;
    st.construct(data, data + 5);
    std::cout << st.query(1, 5) << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5) << '\n';
#elif CHOICE == 1
    // 1. Use primitive type
    // 2. Extra comparator is specified for segment tree
    const int data[] = {1, 5, 2, 3, 4};
    SegmentTree<int, SegTreeOp1> st;
    st.construct(data, data + 5);
    std::cout << st.query(1, 5) << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5) << '\n';
#elif CHOICE == 2
    // The simplest example of using user-defined data structure.
    // At least one of operator< or operator() must be provided.
    SegmentTree<STNode1> st;
    const STNode1 data[] = {1, 5, 2, 3, 4};
    st.construct(data, data + 5);
    std::cout << st.query(1, 5).value << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5).value << '\n';
#elif CHOICE == 3
    // The simplest example of using user-defined data structure.
    // At least one of operator< or operator() must be provided.
    SegmentTree<STNode2> st;
    const STNode2 data[] = {1, 5, 2, 3, 4};
    st.construct(data, data + 5);
    std::cout << st.query(1, 5).value << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5).value << '\n';
#elif CHOICE == 4
    // when both operator< and operator() are provided, SegmentTree will choose operator() for comparison
    SegmentTree<STNode3> st;
    const STNode3 data[] = {1, 5, 2, 3, 4};
    st.construct(data, data + 5);
    std::cout << st.query(1, 5).value << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5).value << '\n';
#elif CHOICE == 5
    // Different from choice 2, the overloaded operator< is non-member function.
    SegmentTree<STNode4> st;
    const STNode4 data[] = {1, 5, 2, 3, 4};
    st.construct(data, data + 5);
    std::cout << st.query(1, 5).value << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5).value << '\n';
#elif CHOICE == 6
    // 1. Use primitive type
    // 2. Extra operator is specified for segment tree
    const int data[] = {1, 5, 2, 3, 4};
    SegmentTree<int, SegTreeOp2> st;
    st.construct(data, data + 5);
    std::cout << st.query(1, 5) << '\n';
    std::cout << st.query(2, 4) << '\n';
    st.update(1, -1);
    std::cout << st.query(1, 5) << '\n';
    std::cout << st.query(2, 4) << '\n';
#else
    std::cout << "choice " << CHOICE << " is not available.\n";
#endif

    return 0;
}