Segment Tree
(src/DataStructure/SegTree.hpp)

• View this file on GitHub
• Last update: 2023-11-12 00:32:16+09:00
• Include: #include "src/DataStructure/SegTree.hpp"

segment tree

モノイドであるものに対し使用できるデータ構造です。
モノイドとはひとつの二項演算と単位元をもつ代数的構造で、

• 結合律 $(a ・ b) ・ c = a ・ (b ・ c)$
• 単位元が存在する ($e ・ a = a ・ e = a$) 以下の2つの条件を満たすものです。

二項演算の計算量が $O(1)$ である時、

• 要素の $1$ 点更新
• ある区間 $[l, r)$ に対して演算を適用した結果の取得

を $O(logN)$ で行うことができます。

二項演算の計算量が $O(f)$ である時は、以下の計算量に $O(f)$ がかかります。

コンストラクタ

1. SegTree<T> (int N, std::function<T(T, T)> operation, T identity)
2. SegTree(const std::vector<T>& vec, std::function<T(T, T)> operation, T identity)

• 型 T
• $T \times T \to T$ の二項演算 operation
• 単位元 identity が必要です。

1.の形式では長さ $N$ の区間として、単位元で初期化されます。
2.の形式では渡された vector の大きさの区間として、 vector の中の値で初期化されます\

計算量

$O(N)$

set(int idx, T value)

idxの場所にvalueを代入する。

計算量

$O(logN)$

get(int idx)

idxの場所の要素を取得する。

計算量

$O(1)$

prod(int l, int r)

$[l, r)$ の区間の全ての要素に対して、コンストラクタに渡した演算を行った結果を返す。

計算量

$O(logN)$

all_prod()

全ての区間に対してコンストラクタに渡した演算を行う。

計算量

$O(1)$

Verified with

• test/DataStructure/SegTree_Point_Add_Range_Sum.test.cpp

Code

#pragma once

#include <cassert>
#include <functional>
#include <iostream>
#include <vector>

namespace lib{
    template <class T>
    class SegTree {
    private:
        const int _N;
        int log, size;
        const std::function<T(T, T)> operation;  // 演算
        const T identity;                        // 単位元
        std::vector<T> node;

        void update(int idx) {
            node[idx] = operation(node[idx << 1], node[(idx << 1) | 1]);
        }

    public:
        SegTree(int N, std::function<T(T, T)> operation, T identity)
            : SegTree(std::vector<T>(N, identity), operation, identity) {}

        SegTree(const std::vector<T>& vec, std::function<T(T, T)> operation,
                T identity)
            : _N(int(vec.size())), operation(operation), identity(identity) {
            int log = 0;
            while ((1U << log) < (unsigned int)(_N)) {
                log++;
            }
            this->log = log;
            this->size = 1 << log;
            node = std::vector<T>(size << 1, identity);
            for (int i = 0; i < _N; i++) {
                node[size + i] = vec[i];
            }
            for (int i = size - 1; i >= 1; i--) {
                update(i);
            }
        }

        void set(int idx, T value) {
            assert(0 <= idx && idx < _N);
            idx += size;
            node[idx] = value;
            for (int i = 1; i <= log; i++) {
                update(idx >> i);
            }
        }

        T get(int idx) {
            assert(0 <= idx && idx < _N);
            return node[idx + size];
        }

        T prod(int l, int r) {
            assert(0 <= l && l <= r && r <= _N);
            if (l == r) {
                return identity;
            }
            T vl = identity, vr = identity;
            l += size;
            r += size;
            while (l < r) {
                if (l & 1) {
                    vl = operation(vl, node[l++]);
                }
                if (r & 1) {
                    vr = operation(node[--r], vr);
                }
                l >>= 1;
                r >>= 1;
            }
            return operation(vl, vr);
        }

        T all_prod() { return node[1]; }

        int max_right(int l, bool (*f)(T)) {
            assert(0 <= l && l <= _N);
            assert(f(identity));
            if (l == _N) return _N;
            l += size;
            T value = identity;
            do {
                while (!(l & 1)) {
                    l >>= 1;
                }
                if (!f(operation(value, node[l]))) {
                    while (l < size) {
                        l <<= 1;
                        if (f(operation(value, node[l]))) {
                            value = operation(value, node[l++]);
                        }
                    }
                    return l - size;
                }
                value = operation(value, node[l++]);
            } while ((l & -l) != l);
            return _N;
        }

        int min_left(int r, bool (*f)(T)) {
            assert(0 <= r && r <= _N);
            assert(f(identity));

            if (r == 0) return 0;
            r += size;
            T value = identity;
            do {
                r--;
                while (r > 1 && !(r & 1)) {
                    r >>= 1;
                }

                if (!f(operation(node[r], value))) {
                    while (r < size) {
                        r = (r << 1) + 1;
                        if (f(operation(node[r], value))) {
                            value = operation(node[r--], value);
                        }
                    }
                    return r + 1 - size;
                }
                value = operation(node[r], value);
            } while ((r & -r) != r);
            return 0;
        }
    };
}

#line 2 "src/DataStructure/SegTree.hpp"

#include <cassert>
#include <functional>
#include <iostream>
#include <vector>

namespace lib{
    template <class T>
    class SegTree {
    private:
        const int _N;
        int log, size;
        const std::function<T(T, T)> operation;  // 演算
        const T identity;                        // 単位元
        std::vector<T> node;

        void update(int idx) {
            node[idx] = operation(node[idx << 1], node[(idx << 1) | 1]);
        }

    public:
        SegTree(int N, std::function<T(T, T)> operation, T identity)
            : SegTree(std::vector<T>(N, identity), operation, identity) {}

        SegTree(const std::vector<T>& vec, std::function<T(T, T)> operation,
                T identity)
            : _N(int(vec.size())), operation(operation), identity(identity) {
            int log = 0;
            while ((1U << log) < (unsigned int)(_N)) {
                log++;
            }
            this->log = log;
            this->size = 1 << log;
            node = std::vector<T>(size << 1, identity);
            for (int i = 0; i < _N; i++) {
                node[size + i] = vec[i];
            }
            for (int i = size - 1; i >= 1; i--) {
                update(i);
            }
        }

        void set(int idx, T value) {
            assert(0 <= idx && idx < _N);
            idx += size;
            node[idx] = value;
            for (int i = 1; i <= log; i++) {
                update(idx >> i);
            }
        }

        T get(int idx) {
            assert(0 <= idx && idx < _N);
            return node[idx + size];
        }

        T prod(int l, int r) {
            assert(0 <= l && l <= r && r <= _N);
            if (l == r) {
                return identity;
            }
            T vl = identity, vr = identity;
            l += size;
            r += size;
            while (l < r) {
                if (l & 1) {
                    vl = operation(vl, node[l++]);
                }
                if (r & 1) {
                    vr = operation(node[--r], vr);
                }
                l >>= 1;
                r >>= 1;
            }
            return operation(vl, vr);
        }

        T all_prod() { return node[1]; }

        int max_right(int l, bool (*f)(T)) {
            assert(0 <= l && l <= _N);
            assert(f(identity));
            if (l == _N) return _N;
            l += size;
            T value = identity;
            do {
                while (!(l & 1)) {
                    l >>= 1;
                }
                if (!f(operation(value, node[l]))) {
                    while (l < size) {
                        l <<= 1;
                        if (f(operation(value, node[l]))) {
                            value = operation(value, node[l++]);
                        }
                    }
                    return l - size;
                }
                value = operation(value, node[l++]);
            } while ((l & -l) != l);
            return _N;
        }

        int min_left(int r, bool (*f)(T)) {
            assert(0 <= r && r <= _N);
            assert(f(identity));

            if (r == 0) return 0;
            r += size;
            T value = identity;
            do {
                r--;
                while (r > 1 && !(r & 1)) {
                    r >>= 1;
                }

                if (!f(operation(node[r], value))) {
                    while (r < size) {
                        r = (r << 1) + 1;
                        if (f(operation(node[r], value))) {
                            value = operation(node[r--], value);
                        }
                    }
                    return r + 1 - size;
                }
                value = operation(node[r], value);
            } while ((r & -r) != r);
            return 0;
        }
    };
}

Back to top page