Binary indexed tree
(include/data_structure/binary_indexed_tree.hpp)

View this file on GitHub
Last update: 2021-08-11 13:38:32+09:00
Include: #include "include/data_structure/binary_indexed_tree.hpp"

数の配列に対して以下のクエリが対数時間で行えるデータ構造である binary_indexed_tree クラスが定義されています。

配列の一つの要素の値を更新する
配列の一つの要素の値を取得する
配列の区間内に含まれる全ての要素に同じ値を加算する
配列の区間内に含まれる要素の和を取得する

クラステンプレート

binary_indexed_tree クラスは要素の型を表す 1 つのテンプレート引数をとります。

lib::binary_indexed_tree<int>  // int 型の要素を持つ binary indexed tree

要素の型には int, long long 等の整数型の他に、double 等の浮動小数点数型や __int128, cpp_int, static_modint 等の加減算の演算子が適切に定義された型を用いることができます。

コンストラクタ

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

lib::binary_indexed_tree<int> tree_1(N);      // 1.
lib::binary_indexed_tree<int> tree_2(N, 10);  // 2.
lib::binary_indexed_tree<int> tree_3(v);      // 3.

要素数 $N$ の、全ての要素が $0$ で初期化された binary indexed tree を構築します。
要素数 $N$ の、全ての要素が $10$ で初期化された binary indexed tree を構築します。
$v$ の要素で初期化された、$v$ と同じ要素数の binary indexed tree を構築します。

メンバ関数

`size()`

要素数を返します。

`add(i, x)`

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

`uniform_add(l, r, x)`

半開区間 $[l, r)$ (0-indexed) に含まれる全ての要素に $x$ を加算します。

`sum(l, r)`

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

`get(i)`

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

`set(i, x)`

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

`debug_print()`

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

Verified with

Code

//! @file binary_indexed_tree.hpp

#ifndef CP_LIBRARY_BINARY_INDEXED_TREE_HPP
#define CP_LIBRARY_BINARY_INDEXED_TREE_HPP

#include <cassert>
#include <iostream>
#include <string>
#include <vector>

#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::binary_indexed_tree_hpp {
  //! @brief Normal binary indexed tree.
  //! @tparam Elem Element type. Watch out for overflows.
  template <typename Elem>
  class binary_indexed_tree_impl {
  private:
    int length;
    std::vector<Elem> data;

    //! @return Sum of the elements within [0, index) (0-indexed, half-open interval)
    [[nodiscard]] Elem partial_sum(int index) const {
      Elem res = 0;
      for (; index > 0; index -= (index & -index))
        res += data[index];
      return res;
    }

  public:
    //! @brief Construct a vector of n zeroes.
    //! @param n vector size
    explicit binary_indexed_tree_impl(const int n) : length(n), data(n + 1, (Elem) 0) {}

    //! @brief Construct a vector from an existing container.
    //! @tparam Container container type (deduced from parameter).
    //! @param src Source (container)
    template <typename Container>
    explicit binary_indexed_tree_impl(const Container& src)
        : length(static_cast<int>(std::size(src))), data(length + 1, (Elem) 0) {
      for (int i = 0; i < length; ++i)
        add(i, src[i]);
    }

    //! @brief Construct a vector of length n filled with initial_values.
    //! @param n vector size
    //! @param initial_value initial value for all elements
    binary_indexed_tree_impl(const int n, const Elem& initial_value) : length(n), data(n + 1, (Elem) 0) {
      for (int i = 0; i < length; ++i)
        add(i, initial_value);
    }

    //! @return Vector 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)
    void add(int index, const Elem& value) {
      CP_LIBRARY_ASSERT(0 <= index && index < length);
      for (++index; index <= length; index += (index & -index))
        data[index] += value;
    }

    //! @brief Calculate interval sum.
    //! @param left lower limit of interval (0-indexed)
    //! @param right upper limit of interval (0-indexed)
    //! @return Sum of the elements within [left, right) (half-open interval)
    //! @note Time complexity: O(log size)
    [[nodiscard]] Elem sum(int left, int right) const {
      CP_LIBRARY_ASSERT(0 <= left && left <= right && right <= length);
      if (left == 0)
        return partial_sum(right);
      else
        return partial_sum(right) - partial_sum(left - 1);
    }

    //! @brief Get the value of the index-th element.
    //! @param index index (0-indexed)
    //! @note Time complexity: O(log size)
    [[nodiscard]] Elem get(int index) const {
      return partial_sum(index + 1) - partial_sum(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)
    void set(const int index, const Elem& value) {
      add(index, value - get(index));
    }
  };
}  // namespace internal::binary_indexed_tree_hpp

//! @brief Binary indexed tree with uniform add function.
//! @tparam Elem Element type. Watch out for overflows.
template <typename Elem>
class binary_indexed_tree {
private:
  internal::binary_indexed_tree_hpp::binary_indexed_tree_impl<Elem> bit_0, bit_1;

  //! @return Sum of the elements within [0, index) (0-indexed, half-open interval)
  [[nodiscard]] Elem partial_sum(int index) const {
    return bit_0.sum(0, index) + bit_1.sum(0, index) * (index - 1);
  }

public:
  //! @brief Construct a vector of n zeroes.
  //! @param n vector size
  explicit binary_indexed_tree(const int n) : bit_0(n), bit_1(n) {}

  //! @brief Construct a vector from an existing container.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  template <typename Container>
  explicit binary_indexed_tree(const Container& src) : bit_0(src), bit_1(static_cast<int>(std::size(src))) {}

  //! @brief Construct a vector of length n filled with initial_values.
  //! @param n vector size
  //! @param initial_value initial value for all elements
  binary_indexed_tree(const int n, const Elem& initial_value) : bit_0(n, initial_value), bit_1(n) {}

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

  //! @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)
  void add(int index, const Elem& value) {
    bit_0.add(index, value);
  }

  //! @brief Add value to the elements within [left, right) (half-open interval)
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @param value value to be added
  //! @note Time complexity: O(log size)
  void uniform_add(int left, int right, const Elem& value) {
    CP_LIBRARY_ASSERT(0 <= left && left <= right && right <= size());
    if (left != size()) {
      bit_0.add(left, value * (-1) * (left - 1));
      bit_1.add(left, value);
    }
    if (right != size()) {
      bit_0.add(right, value * (right - 1));
      bit_1.add(right, value * (-1));
    }
  }

  //! @brief Calculate interval sum.
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @return Sum of the elements within [left, right) (half-open interval)
  //! @note Time complexity: O(log size)
  [[nodiscard]] Elem sum(int left, int right) const {
    if (left == 0)
      return partial_sum(right);
    else
      return partial_sum(right) - partial_sum(left - 1);
  }

  //! @brief Get the value of the index-th element.
  //! @param index index (0-indexed)
  //! @note Time complexity: O(log size)
  [[nodiscard]] Elem get(int index) const {
    return partial_sum(index + 1) - partial_sum(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)
  void set(const int index, const Elem& value) {
    bit_0.add(index, value - get(index));
  }

  //! @brief Print debug information.
  //! @param name variable name
  //! @param os output stream
  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 << "sum [ ";
    for (int i = 0; i <= size(); ++i)
      os << partial_sum(i) << ' ';
    os << "]\n";
#endif
  }
};

}  // namespace lib

#ifdef CP_LIBRARY_ASSERT_NOT_DEFINED
#  undef CP_LIBRARY_ASSERT
#  undef CP_LIBRARY_ASSERT_NOT_DEFINED
#endif

#endif  // CP_LIBRARY_BINARY_INDEXED_TREE_HPP

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

//! @file binary_indexed_tree.hpp

#ifndef CP_LIBRARY_BINARY_INDEXED_TREE_HPP
#define CP_LIBRARY_BINARY_INDEXED_TREE_HPP

#include <cassert>
#include <iostream>
#include <string>
#include <vector>

#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::binary_indexed_tree_hpp {
  //! @brief Normal binary indexed tree.
  //! @tparam Elem Element type. Watch out for overflows.
  template <typename Elem>
  class binary_indexed_tree_impl {
  private:
    int length;
    std::vector<Elem> data;

    //! @return Sum of the elements within [0, index) (0-indexed, half-open interval)
    [[nodiscard]] Elem partial_sum(int index) const {
      Elem res = 0;
      for (; index > 0; index -= (index & -index))
        res += data[index];
      return res;
    }

  public:
    //! @brief Construct a vector of n zeroes.
    //! @param n vector size
    explicit binary_indexed_tree_impl(const int n) : length(n), data(n + 1, (Elem) 0) {}

    //! @brief Construct a vector from an existing container.
    //! @tparam Container container type (deduced from parameter).
    //! @param src Source (container)
    template <typename Container>
    explicit binary_indexed_tree_impl(const Container& src)
        : length(static_cast<int>(std::size(src))), data(length + 1, (Elem) 0) {
      for (int i = 0; i < length; ++i)
        add(i, src[i]);
    }

    //! @brief Construct a vector of length n filled with initial_values.
    //! @param n vector size
    //! @param initial_value initial value for all elements
    binary_indexed_tree_impl(const int n, const Elem& initial_value) : length(n), data(n + 1, (Elem) 0) {
      for (int i = 0; i < length; ++i)
        add(i, initial_value);
    }

    //! @return Vector 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)
    void add(int index, const Elem& value) {
      CP_LIBRARY_ASSERT(0 <= index && index < length);
      for (++index; index <= length; index += (index & -index))
        data[index] += value;
    }

    //! @brief Calculate interval sum.
    //! @param left lower limit of interval (0-indexed)
    //! @param right upper limit of interval (0-indexed)
    //! @return Sum of the elements within [left, right) (half-open interval)
    //! @note Time complexity: O(log size)
    [[nodiscard]] Elem sum(int left, int right) const {
      CP_LIBRARY_ASSERT(0 <= left && left <= right && right <= length);
      if (left == 0)
        return partial_sum(right);
      else
        return partial_sum(right) - partial_sum(left - 1);
    }

    //! @brief Get the value of the index-th element.
    //! @param index index (0-indexed)
    //! @note Time complexity: O(log size)
    [[nodiscard]] Elem get(int index) const {
      return partial_sum(index + 1) - partial_sum(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)
    void set(const int index, const Elem& value) {
      add(index, value - get(index));
    }
  };
}  // namespace internal::binary_indexed_tree_hpp

//! @brief Binary indexed tree with uniform add function.
//! @tparam Elem Element type. Watch out for overflows.
template <typename Elem>
class binary_indexed_tree {
private:
  internal::binary_indexed_tree_hpp::binary_indexed_tree_impl<Elem> bit_0, bit_1;

  //! @return Sum of the elements within [0, index) (0-indexed, half-open interval)
  [[nodiscard]] Elem partial_sum(int index) const {
    return bit_0.sum(0, index) + bit_1.sum(0, index) * (index - 1);
  }

public:
  //! @brief Construct a vector of n zeroes.
  //! @param n vector size
  explicit binary_indexed_tree(const int n) : bit_0(n), bit_1(n) {}

  //! @brief Construct a vector from an existing container.
  //! @tparam Container container type (deduced from parameter).
  //! @param src Source (container)
  template <typename Container>
  explicit binary_indexed_tree(const Container& src) : bit_0(src), bit_1(static_cast<int>(std::size(src))) {}

  //! @brief Construct a vector of length n filled with initial_values.
  //! @param n vector size
  //! @param initial_value initial value for all elements
  binary_indexed_tree(const int n, const Elem& initial_value) : bit_0(n, initial_value), bit_1(n) {}

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

  //! @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)
  void add(int index, const Elem& value) {
    bit_0.add(index, value);
  }

  //! @brief Add value to the elements within [left, right) (half-open interval)
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @param value value to be added
  //! @note Time complexity: O(log size)
  void uniform_add(int left, int right, const Elem& value) {
    CP_LIBRARY_ASSERT(0 <= left && left <= right && right <= size());
    if (left != size()) {
      bit_0.add(left, value * (-1) * (left - 1));
      bit_1.add(left, value);
    }
    if (right != size()) {
      bit_0.add(right, value * (right - 1));
      bit_1.add(right, value * (-1));
    }
  }

  //! @brief Calculate interval sum.
  //! @param left lower limit of interval (0-indexed)
  //! @param right upper limit of interval (0-indexed)
  //! @return Sum of the elements within [left, right) (half-open interval)
  //! @note Time complexity: O(log size)
  [[nodiscard]] Elem sum(int left, int right) const {
    if (left == 0)
      return partial_sum(right);
    else
      return partial_sum(right) - partial_sum(left - 1);
  }

  //! @brief Get the value of the index-th element.
  //! @param index index (0-indexed)
  //! @note Time complexity: O(log size)
  [[nodiscard]] Elem get(int index) const {
    return partial_sum(index + 1) - partial_sum(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)
  void set(const int index, const Elem& value) {
    bit_0.add(index, value - get(index));
  }

  //! @brief Print debug information.
  //! @param name variable name
  //! @param os output stream
  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 << "sum [ ";
    for (int i = 0; i <= size(); ++i)
      os << partial_sum(i) << ' ';
    os << "]\n";
#endif
  }
};

}  // namespace lib

#ifdef CP_LIBRARY_ASSERT_NOT_DEFINED
#  undef CP_LIBRARY_ASSERT
#  undef CP_LIBRARY_ASSERT_NOT_DEFINED
#endif

#endif  // CP_LIBRARY_BINARY_INDEXED_TREE_HPP

Binary indexed tree (include/data_structure/binary_indexed_tree.hpp)

クラステンプレート

コンストラクタ

メンバ関数

size()

add(i, x)

uniform_add(l, r, x)

sum(l, r)

get(i)

set(i, x)

debug_print()