Segment tree
(include/data_structure/segment_tree.hpp)

View this file on GitHub
Last update: 2021-08-14 12:12:32+09:00
Include: #include "include/data_structure/segment_tree.hpp"

配列に対して以下のクエリが対数時間で行えるデータ構造である segment_tree クラスと、その制御に用いる no_range_query_in_this_scope クラスが定義されています。binary_indexed_tree や sparse_table が利用できる場合、そちらを利用した方が高速に動作することが期待できます。

配列の一つの要素の値を更新する
配列の一つの要素の値を取得する
配列の区間積を取得する

ここでいう「区間積」とは、配列のある区間内に二項演算を繰り返し適用した結果です。例えば二項演算を加算とすると、配列 ${a_1,\, a_2,\, a_3,\, a_4,\, a_5,\, a_6}$ の $a_2$ から $a_4$ までの区間の区間積は $a_2 + a_3 + a_4$ です。二項演算は単位元が存在して結合法則を満たすものであれば構いません。例えば掛け算・最小値を求める演算・最大値を求める演算等を使うこともできます。

クラステンプレート

segment_tree クラスは 2 つの省略可能なテンプレート引数をとります。

lib::segment_tree  // 1.
lib::segment_tree<int>  // 2.
lib::segment_tree<lib::static_modint<1000000007>>  // 3.
lib::segment_tree<int, std::bit_xor<>>  // 4.

テンプレート引数を省略した場合、要素の型は long long で二項演算は加算 (+) となります。
第一引数で要素の型を指定できます。この場合要素の型は int で二項演算は加算 (+) となります。
要素の型は組み込み型でなくてもよいです。この場合要素の型は lib::static_modint<1000000007> で二項演算は加算 (+) となります。
第二引数で二項演算を指定できます。この場合要素の型は int で二項演算は bitwise xor (^) となります。後述の通り他にも二項演算を変更する方法があります。

コンストラクタ

std::vector<int> v {1, 2, 3, 4, 5};

lib::segment_tree tree_1(N);      // 1.
lib::segment_tree tree_2(N, 10);  // 2.
lib::segment_tree tree_3(v);      // 3.

lib::segment_tree<long long, std::bit_xor<>> tree_4(N);      // 4.
lib::segment_tree<long long, std::bit_xor<>> tree_5(N, 10);  // 5.
lib::segment_tree<long long, std::bit_xor<>> tree_6(v);      // 6.

// 最小値を返す演算
const auto func_min = [](const long long x, const long long y) constexpr { return std::min(x, y); };
// 最小値を返す演算の単位元
constexpr long long id_min = 1000000000000000005LL;

lib::segment_tree tree_7(N, id_min, func_min);       // 7.
lib::segment_tree tree_8(N, 500, id_min, func_min);  // 8.
lib::segment_tree tree_9(v, id_min, func_min);       // 9.

// 引数の部分にラムダ式を直接書いても良い
// lib::segment_tree tree_9 (v, 1000000000000000005LL,
//                           [](long long x, long long y) { return std::min(x, y); });

要素数 $N$ の、全ての要素が $0$ で初期化された、二項演算を加算とした segment tree を構築します。
要素数 $N$ の、全ての要素が $10$ で初期化された、二項演算を加算とした segment tree を構築します。
$v$ の要素で初期化された、$v$ と同じ要素数の、二項演算を加算とした segment tree を構築します。
要素数 $N$ の、全ての要素が $0$ で初期化された、二項演算を bitwise xor とした segment tree を構築します。
要素数 $N$ の、全ての要素が $10$ で初期化された、二項演算を bitwise xor とした segment tree を構築します。
$v$ の要素で初期化された、$v$ と同じ要素数の、二項演算を bitwise xor とした segment tree を構築します。
要素数 $N$ の、全ての要素が id_min で初期化された、二項演算を最小値をとる演算とした segment tree を構築します。
要素数 $N$ の、全ての要素が $500$ で初期化された、二項演算を最小値をとる演算とした segment tree を構築します。
$v$ の要素で初期化された、$v$ と同じ要素数の、二項演算を最小値をとる演算とした segment tree を構築します。

二項演算として関数へのポインタを渡すこともできます。

unsigned max(unsigned x, unsigned y) {
  return (x >= y) ? x : y;
}

int main() {
  int N;
  std::cin >> N;

  lib::segment_tree tree(N, 0u, &max);  // OK
  //                            ^^^^
}

メンバ関数

`size()`

要素数を返します。

`add(i, x)`

$i$ 番目 (0-indexed) の要素に $x$ を加算します。

`prod(l, r)`

半開区間 $[l, r)$ (0-indexed) に含まれる要素の区間積を返します。

`all_prod(l, r)`

全要素についての区間積を返します (prod(0, size()) と同じ値を返します)。

`get(i)`

$i$ 番目 (0-indexed) の要素の値を返します。

`set(i, x)`

$i$ 番目 (0-indexed) の要素の値を $x$ にします。

`debug_print()`

デバッグレベルが $1$ 以上のとき、標準エラー出力にデバッグ情報(配列の内容)を出力します。

以下は処理をほんのり高速化させるためのメンバ関数です。使い方を間違えると CP_LIBRARY_ASSERT マクロによって異常終了するか、正しくない計算結果が返ります。

これらの関数の使用を検討する場合、他に適したデータ構造が無いかを考えてください。

`lock()`

set 関数や add 関数で要素を変更した時に、その変更が自動的に親ノードに伝播されないようにします。そのため lock() を行った後には (unlock() を行うまでは) 区間積は取得できなくなります。この関数を呼び、unlock() を行わずにもう一度この関数を呼ぶと CP_LIBRARY_ASSERT マクロによって異常終了します。

`unlock()`

set 関数や add 関数で要素を変更した時に、その変更が自動的に親ノードに伝播されるようにします。lock() を呼んでいない状態でこの関数を呼ぶと CP_LIBRARY_ASSERT マクロによって異常終了します。

この関数は内部に保持されている locked という変数の値を false にセットするだけなので、この関数を呼んだだけでは変更が親ノードへ伝播されないことに注意してください。

`propagate(i)`

$i$ 番目の要素の変更を親ノードに伝播します。変更が自動的に親ノードに伝播されない状態で行った変更を手動で反映させるために使います。

`propagate_all()`

全ての要素の変更を親ノードに伝播します。変更が自動的に親ノードに伝播されない状態で行った変更を手動で反映させるために使います。

`propagate_all_and_unlock()`

propagate_all() と unlock() を呼びます。

`is_locked()`

変更が自動的に親ノードに伝播されない状態になっているかどうかを返します。

lock() や unlock() を直接使用するとバグの原因となるので、no_range_query_in_this_scope クラスを通して使用するとよいです。

no_range_query_in_this_scope クラスのコンストラクタで自動的に lock() が呼ばれ、デストラクタで自動的に propagate_all_and_unlock() が呼ばれます。

lib::segmant_tree tree(N);

{
  // コンストラクタが呼ばれ、tree は変更が自動的に親ノードに伝播されない状態になる
  lib::no_range_query_in_this_scope query_lock(tree);

  // 大量の変更クエリ (ここで区間積の取得はできないが、get() を使うことはできる)
  for (int i = 0; i < Q; ++i)
    tree.set(x[i], y[i]);

  // スコープを抜けるのでデストラクタが呼ばれ、tree への変更が適用される
}

Verified with

Code

//! @file segment_tree.hpp

#ifndef CP_LIBRARY_SEGMENT_TREE_HPP
#define CP_LIBRARY_SEGMENT_TREE_HPP

#include <algorithm>
#include <cassert>
#include <functional>
#include <iostream>
#include <limits>
#include <string>
#include <type_traits>
#include <vector>

#ifndef CP_LIBRARY_WARN
#  if (CP_LIBRARY_DEBUG_LEVEL >= 1)
//! @brief Print warning message
//! @note You can suppress the warning by uncommenting the following line
#    define CP_LIBRARY_WARN(msg) (std::cerr << (msg) << '\n')
// #  define CP_LIBRARY_WARN(msg) (static_cast<void>(0))
#  else
#    define CP_LIBRARY_WARN(msg) (static_cast<void>(0))
#  endif
#  define CP_LIBRARY_WARN_NOT_DEFINED
#endif

#ifndef CP_LIBRARY_ASSERT
//! @brief Assert macro
#  define CP_LIBRARY_ASSERT(...) assert(__VA_ARGS__)
#  define CP_LIBRARY_ASSERT_NOT_DEFINED
#endif

namespace lib {

namespace internal::segment_tree_hpp {
  template <typename Func, typename Arg>
  auto is_binary_op_of_impl(int) -> std::bool_constant<std::is_same_v<decltype(std::declval<Func>()(std::declval<Arg>(), std::declval<Arg>())), Arg>>;
  template <typename Func, typename Arg>
  auto is_binary_op_of_impl(...) -> std::false_type;

  //! @brief Check if Func(Arg, Arg) returns a value of type Arg.
  template <typename Func, typename Arg>
  [[maybe_unused]] constexpr bool is_binary_op_of_v = decltype(is_binary_op_of_impl<Func, Arg>(int {}))::value;
}  // namespace internal::segment_tree_hpp

//! @brief Segment tree
//! @tparam Elem element type. Watch out for overflows.
//! @tparam Func binary op type.
template <typename Elem = long long, typename Func = std::plus<>,
          std::enable_if_t<internal::segment_tree_hpp::is_binary_op_of_v<Func, Elem>, std::nullptr_t> = nullptr>
class segment_tree {
private:
  const int length;
  const Elem id;
  std::vector<Elem> data;
  Func binary_op;
  bool locked;

  //! @brief Propagate changes in the index-th element to its ancestors.
  //! @note Time complexity: O(log size)
  void propagate_impl(int index) {
    index += length;
    while (index > 0) {
      index >>= 1;
      data[index] = binary_op(data[index << 1], data[index << 1 | 1]);
    }
  }

public:
  //! @brief Construct a vector of n zeroes. Default operation (sum) will be selected.
  //! @param n vector size
  explicit segment_tree(const int n)
      : length(n),
        id(0),
        data(length << 1, id),
        binary_op(),
        locked(false) {}

  //! @brief Construct a vector from an existing container. Default operation (sum) will be selected.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  template <typename Container>
  explicit segment_tree(const Container& src)
      : length(static_cast<int>(std::size(src))),
        id(0),
        data(length << 1, id),
        binary_op(),
        locked(false) {
    std::copy(std::cbegin(src), std::cend(src), std::begin(data) + length);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of length n filled with init_vals. Default operation (sum) will be selected.
  //! @param n vector size
  //! @param init_val initial value for all elements
  segment_tree(const int n, const Elem init_val)
      : length(n),
        id(0),
        data(length << 1, id),
        binary_op(),
        locked(false) {
    std::fill(std::begin(data) + length, std::end(data), init_val);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of n zeroes, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem identity_element, const Func& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(binary_operation),
        locked(false) {}

  //! @brief Construct a vector from an existing container, specifying a binary operation and its identity element.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  template <typename Container>
  segment_tree(const Container& src, const Elem identity_element, const Func& binary_operation)
      : length(static_cast<int>(std::size(src))),
        id(identity_element),
        data(length << 1, id),
        binary_op(binary_operation),
        locked(false) {
    std::copy(std::cbegin(src), std::cend(src), std::begin(data) + length);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of length n filled with init_vals, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param init_val initial value for all elements
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem init_val, const Elem identity_element, const Func& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(binary_operation),
        locked(false) {
    std::fill(std::begin(data) + length, std::end(data), init_val);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of n zeroes, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem identity_element, Func&& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(std::move(binary_operation)),
        locked(false) {}

  //! @brief Construct a vector from an existing container, specifying a binary operation and its identity element.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  template <typename Container>
  segment_tree(const Container& src, const Elem identity_element, Func&& binary_operation)
      : length(static_cast<int>(std::size(src))),
        id(identity_element),
        data(length << 1, id),
        binary_op(std::move(binary_operation)),
        locked(false) {
    std::copy(std::cbegin(src), std::cend(src), std::begin(data) + length);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of n zeroes, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem init_val, const Elem identity_element, Func&& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(std::move(binary_operation)),
        locked(false) {
    std::fill(std::begin(data) + length, std::end(data), init_val);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  ~segment_tree() {
    if (locked)
      CP_LIBRARY_WARN("Segment tree is destructed in a locked state.");
  }

  //! @return Vector size (length)
  [[nodiscard]] int size() const noexcept {
    return length;
  }

  //! @brief Add value to the index-th element.
  //! @param index index of the element to be added (0-indexed)
  //! @param value value to be added
  //! @note Time complexity: O(log size) if unlocked
  //! @note Time complexity: O(1)        if locked
  void add(const int index, const Elem value) {
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    data[index + length] = data[index + length] + value;
    if (!locked)
      propagate_impl(index);
  }

  //! @brief Set the value of the index-th element to value.
  //! @param index index (0-indexed)
  //! @param value value to be set
  //! @note Time complexity: O(log size) if unlocked
  //! @note Time complexity: O(1)        if locked
  void set(const int index, const Elem value) {
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    data[index + length] = value;
    if (!locked)
      propagate_impl(index);
  }

  //! @brief Get the value of the index-th element.
  //! @param index index (0-indexed)
  //! @note Time complexity: O(1)
  [[nodiscard]] Elem get(const int index) const {
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    return data[index + length];
  }

  //! @brief Calculate interval product.
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @return product (result of the specified binary operation) of the elements within [left, right) (half-open interval)
  //! @note Time complexity: O(log size)
  [[nodiscard]] Elem prod(int L, int R) const {
    CP_LIBRARY_ASSERT(!locked);
    CP_LIBRARY_ASSERT(0 <= L && L <= R && R <= length);
    L += length;
    R += length;

    Elem res_l = id, res_r = id;

    while (L < R) {
      if (L & 1) {
        res_l = binary_op(res_l, data[L]);
        ++L;
      }
      if (R & 1)
        res_r = binary_op(data[--R], res_r);
      L >>= 1;
      R >>= 1;
    }

    return binary_op(res_l, res_r);
  }

  //! @brief Calculate product of all elements.
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @return product (result of the specified binary operation) of all elements
  //! @note Time complexity: O(1)
  [[nodiscard]] Elem all_prod() const {
    CP_LIBRARY_ASSERT(!locked);
    return data[1];
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @brief Propagate changes in the index-th element to its ancestors.
  //! @note Time complexity: O(log size)
  void propagate(int index) {
    CP_LIBRARY_ASSERT(locked);
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    propagate_impl(index);
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @brief Propagate changes of all elements to the ancestors.
  //! @note Time complexity: O(size)
  void propagate_all() {
    CP_LIBRARY_ASSERT(locked);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @brief Propagate changes of all elements to the ancestors and resume automatic propagation.
  //! @note Time complexity: O(size)
  void propagate_all_and_unlock() {
    propagate_all();
    unlock();
  }

  //! @warning You can call this function only if you can promise not to call prod()/all_prod() until you call unlock().
  //! @brief Stop automatic propagation on element changes.
  //! @note Time complexity: O(1)
  void lock() {
    CP_LIBRARY_ASSERT(!locked);
    locked = true;
  }

  //! @warning You can call this function only if this segment tree is locked.
  //! @warning This function will not perform propagation. You need to call propagate()/propagate_all() manually.
  //! @brief Resume automatic propagation on element changes.
  //! @note Time complexity: O(1)
  void unlock() {
    CP_LIBRARY_ASSERT(locked);
    locked = false;
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @return Whether the automatic propagation is stopped.
  //! @note Time complexity: O(1)
  [[nodiscard]] bool is_locked() const {
    return locked;
  }

  void debug_print([[maybe_unused]] const std::string& name = "", [[maybe_unused]] std::ostream& os = std::cerr) const {
#if (CP_LIBRARY_DEBUG_LEVEL >= 1)
    if (!name.empty())
      os << name << ": ";

    os << "val  [ ";
    for (int i = 0; i < size(); ++i)
      os << get(i) << ' ';
    os << "]\n";

    if (!name.empty())
      os << std::string(std::size(name) + 2, ' ');

    os << "prod [ ";

    if (locked)
      os << "cannot display since range product query is locked ";
    else
      for (int i = 0; i <= size(); ++i)
        os << prod(0, i) << ' ';

    os << "]\n";
#endif
  }
};

namespace internal::segment_tree_hpp {
  template <typename Tp>
  [[maybe_unused]] constexpr bool is_segment_tree_v = false;
  template <typename Elem, typename Func>
  [[maybe_unused]] constexpr bool is_segment_tree_v<segment_tree<Elem, Func>> = true;
}  // namespace internal::segment_tree_hpp

//! @brief Utility class to automatically call lock() in constructor and propagate_all_and_unlock() in destructor.
template <typename Tp, std::enable_if_t<internal::segment_tree_hpp::is_segment_tree_v<Tp>, std::nullptr_t> = nullptr>
class no_range_query_in_this_scope {
private:
  Tp& target_segtree;

public:
  //! @brief Lock segment tree by calling lock().
  //! @param segtree target segment tree
  //! @note Time complexity: O(1)
  explicit no_range_query_in_this_scope(Tp& segtree) : target_segtree(segtree) {
    target_segtree.lock();
  }
  //! @brief Unlock segment tree and apply all changes by calling propagate_all_and_unlock().
  //! @note Time complexity: O(size of target segment tree)
  ~no_range_query_in_this_scope() {
    target_segtree.propagate_all_and_unlock();
  }
};

}  // namespace lib

#ifdef CP_LIBRARY_WARN_NOT_DEFINED
#  undef CP_LIBRARY_WARN
#  undef CP_LIBRARY_WARN_NOT_DEFINED
#  ifdef CP_LIBRARY_WARN
#    undef CP_LIBRARY_WARN
#  endif
#endif

#ifdef CP_LIBRARY_ASSERT_NOT_DEFINED
#  undef CP_LIBRARY_ASSERT
#  undef CP_LIBRARY_ASSERT_NOT_DEFINED
#endif

#endif  // CP_LIBRARY_SEGMENT_TREE_HPP

#line 1 "include/data_structure/segment_tree.hpp"

//! @file segment_tree.hpp

#ifndef CP_LIBRARY_SEGMENT_TREE_HPP
#define CP_LIBRARY_SEGMENT_TREE_HPP

#include <algorithm>
#include <cassert>
#include <functional>
#include <iostream>
#include <limits>
#include <string>
#include <type_traits>
#include <vector>

#ifndef CP_LIBRARY_WARN
#  if (CP_LIBRARY_DEBUG_LEVEL >= 1)
//! @brief Print warning message
//! @note You can suppress the warning by uncommenting the following line
#    define CP_LIBRARY_WARN(msg) (std::cerr << (msg) << '\n')
// #  define CP_LIBRARY_WARN(msg) (static_cast<void>(0))
#  else
#    define CP_LIBRARY_WARN(msg) (static_cast<void>(0))
#  endif
#  define CP_LIBRARY_WARN_NOT_DEFINED
#endif

#ifndef CP_LIBRARY_ASSERT
//! @brief Assert macro
#  define CP_LIBRARY_ASSERT(...) assert(__VA_ARGS__)
#  define CP_LIBRARY_ASSERT_NOT_DEFINED
#endif

namespace lib {

namespace internal::segment_tree_hpp {
  template <typename Func, typename Arg>
  auto is_binary_op_of_impl(int) -> std::bool_constant<std::is_same_v<decltype(std::declval<Func>()(std::declval<Arg>(), std::declval<Arg>())), Arg>>;
  template <typename Func, typename Arg>
  auto is_binary_op_of_impl(...) -> std::false_type;

  //! @brief Check if Func(Arg, Arg) returns a value of type Arg.
  template <typename Func, typename Arg>
  [[maybe_unused]] constexpr bool is_binary_op_of_v = decltype(is_binary_op_of_impl<Func, Arg>(int {}))::value;
}  // namespace internal::segment_tree_hpp

//! @brief Segment tree
//! @tparam Elem element type. Watch out for overflows.
//! @tparam Func binary op type.
template <typename Elem = long long, typename Func = std::plus<>,
          std::enable_if_t<internal::segment_tree_hpp::is_binary_op_of_v<Func, Elem>, std::nullptr_t> = nullptr>
class segment_tree {
private:
  const int length;
  const Elem id;
  std::vector<Elem> data;
  Func binary_op;
  bool locked;

  //! @brief Propagate changes in the index-th element to its ancestors.
  //! @note Time complexity: O(log size)
  void propagate_impl(int index) {
    index += length;
    while (index > 0) {
      index >>= 1;
      data[index] = binary_op(data[index << 1], data[index << 1 | 1]);
    }
  }

public:
  //! @brief Construct a vector of n zeroes. Default operation (sum) will be selected.
  //! @param n vector size
  explicit segment_tree(const int n)
      : length(n),
        id(0),
        data(length << 1, id),
        binary_op(),
        locked(false) {}

  //! @brief Construct a vector from an existing container. Default operation (sum) will be selected.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  template <typename Container>
  explicit segment_tree(const Container& src)
      : length(static_cast<int>(std::size(src))),
        id(0),
        data(length << 1, id),
        binary_op(),
        locked(false) {
    std::copy(std::cbegin(src), std::cend(src), std::begin(data) + length);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of length n filled with init_vals. Default operation (sum) will be selected.
  //! @param n vector size
  //! @param init_val initial value for all elements
  segment_tree(const int n, const Elem init_val)
      : length(n),
        id(0),
        data(length << 1, id),
        binary_op(),
        locked(false) {
    std::fill(std::begin(data) + length, std::end(data), init_val);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of n zeroes, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem identity_element, const Func& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(binary_operation),
        locked(false) {}

  //! @brief Construct a vector from an existing container, specifying a binary operation and its identity element.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  template <typename Container>
  segment_tree(const Container& src, const Elem identity_element, const Func& binary_operation)
      : length(static_cast<int>(std::size(src))),
        id(identity_element),
        data(length << 1, id),
        binary_op(binary_operation),
        locked(false) {
    std::copy(std::cbegin(src), std::cend(src), std::begin(data) + length);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of length n filled with init_vals, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param init_val initial value for all elements
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem init_val, const Elem identity_element, const Func& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(binary_operation),
        locked(false) {
    std::fill(std::begin(data) + length, std::end(data), init_val);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of n zeroes, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem identity_element, Func&& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(std::move(binary_operation)),
        locked(false) {}

  //! @brief Construct a vector from an existing container, specifying a binary operation and its identity element.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  template <typename Container>
  segment_tree(const Container& src, const Elem identity_element, Func&& binary_operation)
      : length(static_cast<int>(std::size(src))),
        id(identity_element),
        data(length << 1, id),
        binary_op(std::move(binary_operation)),
        locked(false) {
    std::copy(std::cbegin(src), std::cend(src), std::begin(data) + length);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @brief Construct a vector of n zeroes, specifying a binary operation and its identity element.
  //! @param n vector size
  //! @param identity_element identity element of the binary operation
  //! @param binary_operation associative binary operation (sum, product, min, max, ...)
  segment_tree(const int n, const Elem init_val, const Elem identity_element, Func&& binary_operation)
      : length(n),
        id(identity_element),
        data(length << 1, id),
        binary_op(std::move(binary_operation)),
        locked(false) {
    std::fill(std::begin(data) + length, std::end(data), init_val);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  ~segment_tree() {
    if (locked)
      CP_LIBRARY_WARN("Segment tree is destructed in a locked state.");
  }

  //! @return Vector size (length)
  [[nodiscard]] int size() const noexcept {
    return length;
  }

  //! @brief Add value to the index-th element.
  //! @param index index of the element to be added (0-indexed)
  //! @param value value to be added
  //! @note Time complexity: O(log size) if unlocked
  //! @note Time complexity: O(1)        if locked
  void add(const int index, const Elem value) {
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    data[index + length] = data[index + length] + value;
    if (!locked)
      propagate_impl(index);
  }

  //! @brief Set the value of the index-th element to value.
  //! @param index index (0-indexed)
  //! @param value value to be set
  //! @note Time complexity: O(log size) if unlocked
  //! @note Time complexity: O(1)        if locked
  void set(const int index, const Elem value) {
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    data[index + length] = value;
    if (!locked)
      propagate_impl(index);
  }

  //! @brief Get the value of the index-th element.
  //! @param index index (0-indexed)
  //! @note Time complexity: O(1)
  [[nodiscard]] Elem get(const int index) const {
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    return data[index + length];
  }

  //! @brief Calculate interval product.
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @return product (result of the specified binary operation) of the elements within [left, right) (half-open interval)
  //! @note Time complexity: O(log size)
  [[nodiscard]] Elem prod(int L, int R) const {
    CP_LIBRARY_ASSERT(!locked);
    CP_LIBRARY_ASSERT(0 <= L && L <= R && R <= length);
    L += length;
    R += length;

    Elem res_l = id, res_r = id;

    while (L < R) {
      if (L & 1) {
        res_l = binary_op(res_l, data[L]);
        ++L;
      }
      if (R & 1)
        res_r = binary_op(data[--R], res_r);
      L >>= 1;
      R >>= 1;
    }

    return binary_op(res_l, res_r);
  }

  //! @brief Calculate product of all elements.
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @return product (result of the specified binary operation) of all elements
  //! @note Time complexity: O(1)
  [[nodiscard]] Elem all_prod() const {
    CP_LIBRARY_ASSERT(!locked);
    return data[1];
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @brief Propagate changes in the index-th element to its ancestors.
  //! @note Time complexity: O(log size)
  void propagate(int index) {
    CP_LIBRARY_ASSERT(locked);
    CP_LIBRARY_ASSERT(0 <= index && index < length);
    propagate_impl(index);
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @brief Propagate changes of all elements to the ancestors.
  //! @note Time complexity: O(size)
  void propagate_all() {
    CP_LIBRARY_ASSERT(locked);
    for (int i = length - 1; i > 0; --i)
      data[i] = binary_op(data[i << 1], data[i << 1 | 1]);
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @brief Propagate changes of all elements to the ancestors and resume automatic propagation.
  //! @note Time complexity: O(size)
  void propagate_all_and_unlock() {
    propagate_all();
    unlock();
  }

  //! @warning You can call this function only if you can promise not to call prod()/all_prod() until you call unlock().
  //! @brief Stop automatic propagation on element changes.
  //! @note Time complexity: O(1)
  void lock() {
    CP_LIBRARY_ASSERT(!locked);
    locked = true;
  }

  //! @warning You can call this function only if this segment tree is locked.
  //! @warning This function will not perform propagation. You need to call propagate()/propagate_all() manually.
  //! @brief Resume automatic propagation on element changes.
  //! @note Time complexity: O(1)
  void unlock() {
    CP_LIBRARY_ASSERT(locked);
    locked = false;
  }

  //! @warning You need to call this function ONLY IF you use lock()/unlock() function.
  //! @return Whether the automatic propagation is stopped.
  //! @note Time complexity: O(1)
  [[nodiscard]] bool is_locked() const {
    return locked;
  }

  void debug_print([[maybe_unused]] const std::string& name = "", [[maybe_unused]] std::ostream& os = std::cerr) const {
#if (CP_LIBRARY_DEBUG_LEVEL >= 1)
    if (!name.empty())
      os << name << ": ";

    os << "val  [ ";
    for (int i = 0; i < size(); ++i)
      os << get(i) << ' ';
    os << "]\n";

    if (!name.empty())
      os << std::string(std::size(name) + 2, ' ');

    os << "prod [ ";

    if (locked)
      os << "cannot display since range product query is locked ";
    else
      for (int i = 0; i <= size(); ++i)
        os << prod(0, i) << ' ';

    os << "]\n";
#endif
  }
};

namespace internal::segment_tree_hpp {
  template <typename Tp>
  [[maybe_unused]] constexpr bool is_segment_tree_v = false;
  template <typename Elem, typename Func>
  [[maybe_unused]] constexpr bool is_segment_tree_v<segment_tree<Elem, Func>> = true;
}  // namespace internal::segment_tree_hpp

//! @brief Utility class to automatically call lock() in constructor and propagate_all_and_unlock() in destructor.
template <typename Tp, std::enable_if_t<internal::segment_tree_hpp::is_segment_tree_v<Tp>, std::nullptr_t> = nullptr>
class no_range_query_in_this_scope {
private:
  Tp& target_segtree;

public:
  //! @brief Lock segment tree by calling lock().
  //! @param segtree target segment tree
  //! @note Time complexity: O(1)
  explicit no_range_query_in_this_scope(Tp& segtree) : target_segtree(segtree) {
    target_segtree.lock();
  }
  //! @brief Unlock segment tree and apply all changes by calling propagate_all_and_unlock().
  //! @note Time complexity: O(size of target segment tree)
  ~no_range_query_in_this_scope() {
    target_segtree.propagate_all_and_unlock();
  }
};

}  // namespace lib

#ifdef CP_LIBRARY_WARN_NOT_DEFINED
#  undef CP_LIBRARY_WARN
#  undef CP_LIBRARY_WARN_NOT_DEFINED
#  ifdef CP_LIBRARY_WARN
#    undef CP_LIBRARY_WARN
#  endif
#endif

#ifdef CP_LIBRARY_ASSERT_NOT_DEFINED
#  undef CP_LIBRARY_ASSERT
#  undef CP_LIBRARY_ASSERT_NOT_DEFINED
#endif

#endif  // CP_LIBRARY_SEGMENT_TREE_HPP

Segment tree (include/data_structure/segment_tree.hpp)

クラステンプレート

コンストラクタ

メンバ関数

size()

add(i, x)

prod(l, r)

all_prod(l, r)

get(i)

set(i, x)

debug_print()

lock()

unlock()

propagate(i)

propagate_all()

propagate_all_and_unlock()

is_locked()