Segment Tree
(DataStructure/SegmentTree.cpp)

• View this file on GitHub
• Last update: 2021-11-21 13:46:02+09:00

説明

配列に対する 1 点更新クエリと区間求値クエリを高速に処理できる。

• SegmentTree<T> seg(int size, T e, std::function<T(T, T)> operation, std::function<T(T, T)> update) : 以下のような配列 A を構成する。計算量 $O(size)$
  • size : 配列の長さ
  • e : Tの単位元かつ A の全ての要素の初期値
  • operation : 求値クエリで使う二項演算
  • update : 更新クエリで使う二項演算
• SegmentTree<T> seg(std::vector<T> v, T e, std::function<T(T, T)> operation, std::function<T(T, T)> : 配列 v で初期化したものを構成する。計算量 $O(size)$
• void seg.set_val(int i, T x) : A[i] = update(A[i], x) の操作を行う。計算量 $O(\log size)$
• T seg.fold(int l, int r) : operation(A[l], A[l + 1], …, A[r - 1]) を求める。計算量 $O(\log size)$
• T seg[i] : A[i] を返す。計算量 $O(1)$
• int seg.max_right(int l, std::function<bool(T)> f) : f(A[l], A[l + 1], …, A[r]) == true を満たす最大の r を求める。計算量 $O(\log size)$

注意 上に書かれている計算量は operation 、 update が定数時間で動くものと仮定して書かれている。これらの内部の計算量が $O(f(n))$ だった場合、全ての計算量が $O(f(n))$ 倍される。

Verified with

• Test/AOJ/SegmentTree-RangeMinimumQuery.test.cpp
• Test/AOJ/SegmentTree-RangeSumQuery.test.cpp
• Test/yosupo-judge/SegmentTree-PointAddRangeSum.test.cpp

Code

template<typename T>
class SegmentTree {
    int internal_size, seg_size;
    const T def;
    std::vector<T> dat;
    std::function<T(T, T)> operation;
    std::function<T(T, T)> update;

public:
    explicit SegmentTree(const int size, const T& e, std::function<T(T, T)> operation, std::function<T(T, T)> update): internal_size(size), def(e), operation(operation), update(update) {
        for(seg_size = 1; seg_size < internal_size; seg_size *= 2);
        dat = std::vector<T>(2 * seg_size, def);
    }

    explicit SegmentTree(std::vector<T> v, const T& e,std::function<T(T, T)> operation, std::function<T(T, T)> update): internal_size(v.size()), def(e), operation(operation), update(update) {
        for(seg_size = 1; seg_size < internal_size; seg_size *= 2);
        dat.resize(2 * seg_size, def);

        for(int i = 0; i < seg_size;i++) dat[i + seg_size] = v[i];
        for(int i = seg_size - 1; i >= 1; i--) {
            dat[i] = operation(dat[i * 2], dat[i * 2 + 1]);
        }
    }

    void set_val(int i, const T& value) {
        i += seg_size;
        dat[i] = update(dat[i], value);
        while(i > 1) {
            i >>= 1;
            dat[i] = operation(dat[i * 2], dat[i * 2 + 1]);
        }
    }

    T fold(int l,int r) {
        l += seg_size;
        r += seg_size;
        T ret = def;
        while(l < r) {
            if(l & 1) ret = operation(ret,dat[l++]);
            if(r & 1) ret = operation(dat[--r],ret);
            l >>= 1;
            r >>= 1;
        }
        return ret;
    }

    int max_right(int l, std::function<bool(T)> f) {
        if(l == internal_size) return internal_size;
        l += seg_size;
        T sum = def;
        do {
            while(!(l & 1)) l >>= 1;
            if(!f(operation(sum, dat[l]))) {
                while(l < seg_size) {
                    l <<= 1;
                    if(f(operation(sum, dat[l]))) sum = operation(sum, dat[l++]);
                }
                return l - seg_size;
            }
        } while((l & -1) != l);
        return internal_size;
    }

    T operator[](int i) {return dat[i + seg_size];}
};

#line 1 "DataStructure/SegmentTree.cpp"
template<typename T>
class SegmentTree {
    int internal_size, seg_size;
    const T def;
    std::vector<T> dat;
    std::function<T(T, T)> operation;
    std::function<T(T, T)> update;

public:
    explicit SegmentTree(const int size, const T& e, std::function<T(T, T)> operation, std::function<T(T, T)> update): internal_size(size), def(e), operation(operation), update(update) {
        for(seg_size = 1; seg_size < internal_size; seg_size *= 2);
        dat = std::vector<T>(2 * seg_size, def);
    }

    explicit SegmentTree(std::vector<T> v, const T& e,std::function<T(T, T)> operation, std::function<T(T, T)> update): internal_size(v.size()), def(e), operation(operation), update(update) {
        for(seg_size = 1; seg_size < internal_size; seg_size *= 2);
        dat.resize(2 * seg_size, def);

        for(int i = 0; i < seg_size;i++) dat[i + seg_size] = v[i];
        for(int i = seg_size - 1; i >= 1; i--) {
            dat[i] = operation(dat[i * 2], dat[i * 2 + 1]);
        }
    }

    void set_val(int i, const T& value) {
        i += seg_size;
        dat[i] = update(dat[i], value);
        while(i > 1) {
            i >>= 1;
            dat[i] = operation(dat[i * 2], dat[i * 2 + 1]);
        }
    }

    T fold(int l,int r) {
        l += seg_size;
        r += seg_size;
        T ret = def;
        while(l < r) {
            if(l & 1) ret = operation(ret,dat[l++]);
            if(r & 1) ret = operation(dat[--r],ret);
            l >>= 1;
            r >>= 1;
        }
        return ret;
    }

    int max_right(int l, std::function<bool(T)> f) {
        if(l == internal_size) return internal_size;
        l += seg_size;
        T sum = def;
        do {
            while(!(l & 1)) l >>= 1;
            if(!f(operation(sum, dat[l]))) {
                while(l < seg_size) {
                    l <<= 1;
                    if(f(operation(sum, dat[l]))) sum = operation(sum, dat[l++]);
                }
                return l - seg_size;
            }
        } while((l & -1) != l);
        return internal_size;
    }

    T operator[](int i) {return dat[i + seg_size];}
};

Back to top page