Segment Tree (セグメント木) (structure/segment-tree/segment-tree.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2024-08-10 03:11:07+09:00
• Include: #include "structure/segment-tree/segment-tree.hpp"

セグメント木は、モノイドについて区間に対する演算が $O(\log N)$ で処理できます。

モノイドは次の条件を満たす代数的構造です。

• 結合律を満たします。つまり $S$ の各元 $a, b, c$ に対して、$(a \cdot b) \cdot c = a \cdot (b \cdot c)$ を満たします。
• 単位元を持ちます。つまり $S$ の任意の元 $a$ をとってきたときに $a \cdot e = e \cdot a = a$ なる $e$ が存在します。

以下の実装では木を 1-indexed の配列で表現しています。ノード $k$ について、親ノードは $\frac k 2$, 子ノードは $2k$, $2k+1$ です。

コンストラクタ

(1) SegmentTree< Monoid >(Monoid m, int n)
(2) SegmentTree< Monoid >(Monoid m, const vector<S> &v) 

1. モノイド m、サイズ n で初期化します。各要素には単位元 m.e() が格納されます。
2. モノイド m、配列 v で初期化します。

計算量

• $O(n)$

Monoid について

Monoid は、次の構造体と関数を持つ構造体です。

struct Monoid {
  using S = ?;
  static constexpr S op(const S& a, const S& b) {}
  static constexpr S e() {}
};

• 型 S
• 二項演算 S op(S a, S b)
• 単位元 e()

例えば、区間和を求めたい場合は次の構造体を定義して、SegmentTree に渡します。

struct RangeSum {
  using S = int;
  static constexpr S op(const S& a, const S& b) { return a + b; }
  static constexpr S e() { return 0; }
};

SegmentTree seg(RangeSum(), n);

LambdaMonoid について

LambdaMonoid は、ラムダ式を受け取って、構造体 Monoid のようにふるまう構造体です。LambdaMonoid の引数に二項演算 S op(S a, S b)、単位元 e() の順で渡すことで初期化できます。

template< typename Op, typename E >
LambdaMonoid(Op _op, E _e)

例えば、区間和を求めたい場合はラムダ式を使って次のように書くこともできます。

auto op = [](int a, int b) { return a + b; };
auto e = []() { return 0; };
SegmentTree seg(LambdaMonoid(op, e), n);

build

void build(const vector<S> &v)

配列 v で初期化します。

制約

• コンストラクタで渡した n と v の長さが一致する

計算量

• $O(n)$

set

void set(int k, const S &x)

k 番目の要素を x に変更します。

制約

• $0 \leq k \lt n$

計算量

• $O(\log n)$

get

S get(int k) const

k 番目の要素を返します。

制約

• $0 \leq k \lt n$

計算量

• $O(1)$

operator[]

S operator[](int k) const

k 番目の要素を返します。

制約

• $0 \leq k \lt n$

計算量

• $O(1)$

apply

void apply(int k, const S &x)

k 番目の要素を、その要素と x を二項演算した値に変更します。

制約

• $0 \leq k \lt n$

計算量

• $O(\log n)$

prod

S prod(int l, int r) const

区間 $[l, r)$ に対して二項演算した結果を返します。

制約

• $0 \leq l \leq r \leq n$

計算量

• $O(\log n)$

all_prod

S all_prod() const

すべての要素を二項演算した結果を返します。

計算量

• $O(1)$

find_first

template <typename C>
int find_first(int l, const C &check) const 

$[a, x)$ が check を満たす最初の要素位置 $x$ を返します。存在しないとき $n$ を返します。

制約

• $0 \leq l \leq n$

計算量

• $O(\log n)$

find_last

template <typename C>
int find_last(int r, const C &check) const

$[x, b)$ が check を満たす最後の要素位置 $x$ を返します。存在しないとき $-1$ を返します。

制約

• $0 \leq r \leq n$

計算量

• $O(\log n)$

Depends on

• structure/class/monoid.hpp

Verified with

• test/verify/aoj-dsl-2-a.test.cpp
• test/verify/yosupo-point-set-range-composite-3.test.cpp
• test/verify/yukicoder-650.test.cpp

Code

#include "../class/monoid.hpp"

template <typename Monoid>
struct SegmentTree {
  using S = typename Monoid::S;

 private:
  int n, sz;

  vector<S> seg;

  Monoid m;

 public:
  SegmentTree() = default;

  explicit SegmentTree(Monoid m, int n) : m(m), n(n) {
    sz = 1;
    while (sz < n) sz <<= 1;
    seg.assign(2 * sz, m.e());
  }

  explicit SegmentTree(Monoid m, const vector<S> &v)
      : SegmentTree(m, (int)v.size()) {
    build(v);
  }

  void build(const vector<S> &v) {
    assert(n == (int)v.size());
    for (int k = 0; k < n; k++) seg[k + sz] = v[k];
    for (int k = sz - 1; k > 0; k--) {
      seg[k] = m.op(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  void set(int k, const S &x) {
    k += sz;
    seg[k] = x;
    while (k >>= 1) {
      seg[k] = m.op(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  S get(int k) const { return seg[k + sz]; }

  S operator[](int k) const { return get(k); }

  void apply(int k, const S &x) {
    k += sz;
    seg[k] = m.op(seg[k], x);
    while (k >>= 1) {
      seg[k] = m.op(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  S prod(int l, int r) const {
    if (l >= r) return m.e();
    S L = m.e(), R = m.e();
    for (l += sz, r += sz; l < r; l >>= 1, r >>= 1) {
      if (l & 1) L = m.op(L, seg[l++]);
      if (r & 1) R = m.op(seg[--r], R);
    }
    return m.op(L, R);
  }

  S all_prod() const { return seg[1]; }

  template <typename C>
  int find_first(int l, const C &check) const {
    if (l >= n) return n;
    l += sz;
    S sum = m.e();
    do {
      while ((l & 1) == 0) l >>= 1;
      if (check(m.op(sum, seg[l]))) {
        while (l < sz) {
          l <<= 1;
          auto nxt = m.op(sum, seg[l]);
          if (not check(nxt)) {
            sum = nxt;
            l++;
          }
        }
        return l + 1 - sz;
      }
      sum = m.op(sum, seg[l++]);
    } while ((l & -l) != l);
    return n;
  }

  template <typename C>
  int find_last(int r, const C &check) const {
    if (r <= 0) return -1;
    r += sz;
    S sum = m.e();
    do {
      r--;
      while (r > 1 and (r & 1)) r >>= 1;
      if (check(m.op(seg[r], sum))) {
        while (r < sz) {
          r = (r << 1) + 1;
          auto nxt = m.op(seg[r], sum);
          if (not check(nxt)) {
            sum = nxt;
            r--;
          }
        }
        return r - sz;
      }
      sum = m.op(seg[r], sum);
    } while ((r & -r) != r);
    return -1;
  }
};

#line 2 "structure/class/monoid.hpp"

template <typename S2, typename Op, typename E>
struct LambdaMonoid {
  using S = S2;
  S op(const S &a, const S &b) const { return _op(a, b); }

  S e() const { return _e(); }

  LambdaMonoid(Op _op, E _e) : _op(_op), _e(_e) {}

 private:
  Op _op;

  E _e;
};

template <typename Op, typename E>
LambdaMonoid(Op _op, E _e) -> LambdaMonoid<decltype(_e()), Op, E>;

/*
struct Monoid {
  using S = ?;
  static constexpr S op(const S& a, const S& b) {}
  static constexpr S e() {}
};
*/
#line 2 "structure/segment-tree/segment-tree.hpp"

template <typename Monoid>
struct SegmentTree {
  using S = typename Monoid::S;

 private:
  int n, sz;

  vector<S> seg;

  Monoid m;

 public:
  SegmentTree() = default;

  explicit SegmentTree(Monoid m, int n) : m(m), n(n) {
    sz = 1;
    while (sz < n) sz <<= 1;
    seg.assign(2 * sz, m.e());
  }

  explicit SegmentTree(Monoid m, const vector<S> &v)
      : SegmentTree(m, (int)v.size()) {
    build(v);
  }

  void build(const vector<S> &v) {
    assert(n == (int)v.size());
    for (int k = 0; k < n; k++) seg[k + sz] = v[k];
    for (int k = sz - 1; k > 0; k--) {
      seg[k] = m.op(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  void set(int k, const S &x) {
    k += sz;
    seg[k] = x;
    while (k >>= 1) {
      seg[k] = m.op(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  S get(int k) const { return seg[k + sz]; }

  S operator[](int k) const { return get(k); }

  void apply(int k, const S &x) {
    k += sz;
    seg[k] = m.op(seg[k], x);
    while (k >>= 1) {
      seg[k] = m.op(seg[2 * k + 0], seg[2 * k + 1]);
    }
  }

  S prod(int l, int r) const {
    if (l >= r) return m.e();
    S L = m.e(), R = m.e();
    for (l += sz, r += sz; l < r; l >>= 1, r >>= 1) {
      if (l & 1) L = m.op(L, seg[l++]);
      if (r & 1) R = m.op(seg[--r], R);
    }
    return m.op(L, R);
  }

  S all_prod() const { return seg[1]; }

  template <typename C>
  int find_first(int l, const C &check) const {
    if (l >= n) return n;
    l += sz;
    S sum = m.e();
    do {
      while ((l & 1) == 0) l >>= 1;
      if (check(m.op(sum, seg[l]))) {
        while (l < sz) {
          l <<= 1;
          auto nxt = m.op(sum, seg[l]);
          if (not check(nxt)) {
            sum = nxt;
            l++;
          }
        }
        return l + 1 - sz;
      }
      sum = m.op(sum, seg[l++]);
    } while ((l & -l) != l);
    return n;
  }

  template <typename C>
  int find_last(int r, const C &check) const {
    if (r <= 0) return -1;
    r += sz;
    S sum = m.e();
    do {
      r--;
      while (r > 1 and (r & 1)) r >>= 1;
      if (check(m.op(seg[r], sum))) {
        while (r < sz) {
          r = (r << 1) + 1;
          auto nxt = m.op(seg[r], sum);
          if (not check(nxt)) {
            sum = nxt;
            r--;
          }
        }
        return r - sz;
      }
      sum = m.op(seg[r], sum);
    } while ((r & -r) != r);
    return -1;
  }
};

Back to top page