Persistent Segment Tree (永続セグメント木) (structure/segment-tree/persistent-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/persistent-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$ が存在します。

永続セグメント木は、セグメント木を永続にしたデータ構造です。

コンストラクタ

PersistentSegmentTree< Monoid >(Monoid m, int n)

モノイド m、サイズ n で初期化します。

計算量

$O(1)$

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; }
};

PersistentSegmentTree 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; };
PersistentSegmentTree seg(LambdaMonoid(op, e), n);

build

NP build(const vector<S> &v) const

配列 v で初期化して、セグメント木のポインタを返します。

制約

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

計算量

$O(n)$

set

NP set(NP t, int k, const S &x) const

セグメント木 t の k 番目の要素を x に変更して、新しいセグメント木へのポインタを返します。

制約

$0 \leq k \lt n$

計算量

$O(\log n)$

get

S get(NP t, int k) const

セグメント木 t の k 番目の要素を返します。

制約

$0 \leq k \lt n$

計算量

$O(1)$

apply

NP apply(NP t, int k, const S &x)

セグメント木 t の k 番目の要素を、その要素と x を二項演算した値に変更し、新しいセグメント木へのポインタを返します。

制約

$0 \leq k \lt n$

計算量

$O(\log n)$

prod

S prod(NP t, int l, int r) const

セグメント木 t の区間 $[l, r)$ に対して二項演算した結果を返します。

制約

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

計算量

$O(\log n)$

all_prod

S all_prod(NP t) const

セグメント木 t のすべての要素を二項演算した結果を返します。

計算量

$O(1)$

Depends on

structure/class/monoid.hpp

Verified with

test/verify/yosupo-range-kth-smallest-3.test.cpp

Code

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

template <typename Monoid>
struct PersistentSegmentTree {
  using S = typename Monoid::S;
  struct Node {
    S d;
    Node *l, *r;
  };
  using NP = Node *;

 private:
  int n{};
  Monoid m;

  NP merge(NP l, NP r) const { return new Node{m.op(l->d, r->d), l, r}; }

  NP build(int l, int r, const vector<S> &v) const {
    if (l + 1 == r) return new Node{v[l], nullptr, nullptr};
    NP lp = build(l, (l + r) / 2, v);
    NP rp = build((l + r) / 2, r, v);
    return merge(lp, rp);
  }

  NP set(int a, const S &x, NP k, int l, int r) const {
    if (r <= a || a + 1 <= l) {
      return k;
    } else if (a <= l && r <= a + 1) {
      return new Node{x, nullptr, nullptr};
    } else {
      return merge(set(a, x, k->l, l, (l + r) >> 1),
                   set(a, x, k->r, (l + r) >> 1, r));
    }
  }

  NP apply(int a, const S &x, NP k, int l, int r) const {
    if (r <= a || a + 1 <= l) {
      return k;
    } else if (a <= l && r <= a + 1) {
      return new Node{m.op(k->d, x), nullptr, nullptr};
    } else {
      return merge(apply(a, x, k->l, l, (l + r) >> 1),
                   apply(a, x, k->r, (l + r) >> 1, r));
    }
  }

  S prod(int a, int b, NP k, int l, int r) const {
    if (r <= a || b <= l) {
      return m.e();
    } else if (a <= l && r <= b) {
      return k->d;
    } else {
      return m.op(prod(a, b, k->l, l, (l + r) >> 1),
                  prod(a, b, k->r, (l + r) >> 1, r));
    }
  }

 public:
  PersistentSegmentTree() = default;

  explicit PersistentSegmentTree(Monoid m, int n) : m(m), n(n) {}

  NP build(const vector<S> &v) const {
    assert(n == (int)v.size());
    return build(0, n, v);
  }

  NP set(NP t, int k, const S &x) const { return set(k, x, t, 0, n); }

  S get(NP t, int k) const {
    int l = 0, r = n;
    while (l + 1 < r) {
      int p = (l + r) / 2;
      if (k < p) {
        t = t->l;
        r = p;
      } else {
        t = t->r;
        l = p;
      }
    }
    return t->d;
  }

  NP apply(NP t, int k, const S &x) const { return apply(k, x, t, 0, n); }

  S prod(NP t, int a, int b) const { return prod(a, b, t, 0, n); }

  S all_prod(NP t) const { return t->d; }
};

#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/persistent-segment-tree.hpp"

template <typename Monoid>
struct PersistentSegmentTree {
  using S = typename Monoid::S;
  struct Node {
    S d;
    Node *l, *r;
  };
  using NP = Node *;

 private:
  int n{};
  Monoid m;

  NP merge(NP l, NP r) const { return new Node{m.op(l->d, r->d), l, r}; }

  NP build(int l, int r, const vector<S> &v) const {
    if (l + 1 == r) return new Node{v[l], nullptr, nullptr};
    NP lp = build(l, (l + r) / 2, v);
    NP rp = build((l + r) / 2, r, v);
    return merge(lp, rp);
  }

  NP set(int a, const S &x, NP k, int l, int r) const {
    if (r <= a || a + 1 <= l) {
      return k;
    } else if (a <= l && r <= a + 1) {
      return new Node{x, nullptr, nullptr};
    } else {
      return merge(set(a, x, k->l, l, (l + r) >> 1),
                   set(a, x, k->r, (l + r) >> 1, r));
    }
  }

  NP apply(int a, const S &x, NP k, int l, int r) const {
    if (r <= a || a + 1 <= l) {
      return k;
    } else if (a <= l && r <= a + 1) {
      return new Node{m.op(k->d, x), nullptr, nullptr};
    } else {
      return merge(apply(a, x, k->l, l, (l + r) >> 1),
                   apply(a, x, k->r, (l + r) >> 1, r));
    }
  }

  S prod(int a, int b, NP k, int l, int r) const {
    if (r <= a || b <= l) {
      return m.e();
    } else if (a <= l && r <= b) {
      return k->d;
    } else {
      return m.op(prod(a, b, k->l, l, (l + r) >> 1),
                  prod(a, b, k->r, (l + r) >> 1, r));
    }
  }

 public:
  PersistentSegmentTree() = default;

  explicit PersistentSegmentTree(Monoid m, int n) : m(m), n(n) {}

  NP build(const vector<S> &v) const {
    assert(n == (int)v.size());
    return build(0, n, v);
  }

  NP set(NP t, int k, const S &x) const { return set(k, x, t, 0, n); }

  S get(NP t, int k) const {
    int l = 0, r = n;
    while (l + 1 < r) {
      int p = (l + r) / 2;
      if (k < p) {
        t = t->l;
        r = p;
      } else {
        t = t->r;
        l = p;
      }
    }
    return t->d;
  }

  NP apply(NP t, int k, const S &x) const { return apply(k, x, t, 0, n); }

  S prod(NP t, int a, int b) const { return prod(a, b, t, 0, n); }

  S all_prod(NP t) const { return t->d; }
};