segment_tree/segment_tree.hpp

• View this file on GitHub
• Last update: 2023-07-03 22:10:06+09:00
• Include: #include "segment_tree/segment_tree.hpp"

Verified with

• test/aoj/0646.test.cpp
• test/aoj/2444_1.test.cpp
• test/aoj/DSL2A_1.test.cpp
• test/aoj/DSL2A_2.test.cpp
• test/aoj/DSL2B.test.cpp

Code

#ifndef SEGMENT_TREE_HPP
#define SEGMENT_TREE_HPP

#include <functional>
#include <iterator>
#include <vector>

//===
template <class Monoid>
struct SegmentTree {
    using T = typename Monoid::value_type;

    std::vector<T> tree;

    SegmentTree() = default;
    explicit SegmentTree(int n) : tree(n << 1, Monoid::identity()){};

    template <class InputIterator>
    SegmentTree(InputIterator first, InputIterator last) {
        tree.assign(distance(first, last) << 1, Monoid::identity());

        int i;
        i = size();
        for (InputIterator itr = first; itr != last; itr++) {
            tree[i++] = *itr;
        }
        for (i = size() - 1; i > 0; i--) {
            tree[i] = Monoid::operation(tree[(i << 1)], tree[(i << 1) | 1]);
        }
    };

    inline int size() {
        return tree.size() >> 1;
    };

    inline T operator[](int k) {
        return tree[k + size()];
    };

    void set(int k, const T dat) {
        k += size();
        tree[k] = dat;

        while (k > 1) {
            k >>= 1;
            tree[k] = Monoid::operation(tree[(k << 1)], tree[(k << 1) | 1]);
        }
    };

    void update(int k, const T dat) {
        set(k, dat);
    };

    // [l, r)
    T fold(int l, int r) {
        l += size();  // points leaf
        r += size();

        T lv = Monoid::identity();
        T rv = Monoid::identity();
        while (l < r) {
            if (l & 1) lv = Monoid::operation(lv, tree[l++]);
            if (r & 1) rv = Monoid::operation(tree[--r], rv);
            l >>= 1;
            r >>= 1;
        }

        return Monoid::operation(lv, rv);
    };

    template <class F>
    inline int sub_tree_search(int i, T sum, F f) {
        while (i < size()) {
            T x = Monoid::operation(sum, tree[i << 1]);
            if (f(x)) {
                i = i << 1;
            } else {
                sum = x;
                i = (i << 1) | 1;
            }
        }
        return i - size() + 1;
    }

    template <class F>
    int find_first(int l, F f) {
        l += size();
        int r = size() * 2;  // r = n;
        int tmpr = r;
        int shift = 0;

        T sum = Monoid::identity();
        while (l < r) {
            if (l & 1) {
                if (f(Monoid::operation(sum, tree[l]))) {
                    return sub_tree_search(l, sum, f);
                }
                sum = Monoid::operation(sum, tree[l]);
                l++;
            }
            l >>= 1;
            r >>= 1;
            shift++;
        }

        while (shift > 0) {
            shift--;
            r = tmpr >> shift;
            if (r & 1) {
                r--;
                if (f(Monoid::operation(sum, tree[r]))) {
                    return sub_tree_search(r, sum, f);
                }
                sum = Monoid::operation(sum, tree[r]);
            }
        }

        return -1;
    };
};
//===

#endif

#line 1 "segment_tree/segment_tree.hpp"



#include <functional>
#include <iterator>
#include <vector>

//===
template <class Monoid>
struct SegmentTree {
    using T = typename Monoid::value_type;

    std::vector<T> tree;

    SegmentTree() = default;
    explicit SegmentTree(int n) : tree(n << 1, Monoid::identity()){};

    template <class InputIterator>
    SegmentTree(InputIterator first, InputIterator last) {
        tree.assign(distance(first, last) << 1, Monoid::identity());

        int i;
        i = size();
        for (InputIterator itr = first; itr != last; itr++) {
            tree[i++] = *itr;
        }
        for (i = size() - 1; i > 0; i--) {
            tree[i] = Monoid::operation(tree[(i << 1)], tree[(i << 1) | 1]);
        }
    };

    inline int size() {
        return tree.size() >> 1;
    };

    inline T operator[](int k) {
        return tree[k + size()];
    };

    void set(int k, const T dat) {
        k += size();
        tree[k] = dat;

        while (k > 1) {
            k >>= 1;
            tree[k] = Monoid::operation(tree[(k << 1)], tree[(k << 1) | 1]);
        }
    };

    void update(int k, const T dat) {
        set(k, dat);
    };

    // [l, r)
    T fold(int l, int r) {
        l += size();  // points leaf
        r += size();

        T lv = Monoid::identity();
        T rv = Monoid::identity();
        while (l < r) {
            if (l & 1) lv = Monoid::operation(lv, tree[l++]);
            if (r & 1) rv = Monoid::operation(tree[--r], rv);
            l >>= 1;
            r >>= 1;
        }

        return Monoid::operation(lv, rv);
    };

    template <class F>
    inline int sub_tree_search(int i, T sum, F f) {
        while (i < size()) {
            T x = Monoid::operation(sum, tree[i << 1]);
            if (f(x)) {
                i = i << 1;
            } else {
                sum = x;
                i = (i << 1) | 1;
            }
        }
        return i - size() + 1;
    }

    template <class F>
    int find_first(int l, F f) {
        l += size();
        int r = size() * 2;  // r = n;
        int tmpr = r;
        int shift = 0;

        T sum = Monoid::identity();
        while (l < r) {
            if (l & 1) {
                if (f(Monoid::operation(sum, tree[l]))) {
                    return sub_tree_search(l, sum, f);
                }
                sum = Monoid::operation(sum, tree[l]);
                l++;
            }
            l >>= 1;
            r >>= 1;
            shift++;
        }

        while (shift > 0) {
            shift--;
            r = tmpr >> shift;
            if (r & 1) {
                r--;
                if (f(Monoid::operation(sum, tree[r]))) {
                    return sub_tree_search(r, sum, f);
                }
                sum = Monoid::operation(sum, tree[r]);
            }
        }

        return -1;
    };
};
//===

Back to top page