任意次元Fenwick Treeとその実装

$$ \gdef\defarrow{\overset{\text{def}}{\Longleftrightarrow}} $$

注: 計算量解析が全体的に怪しいです。疑いながら読んでください(?)

問題

$D$ 次元空間の格子点 $\bm{x}\in \Z ^ D$ の重み $w(\bm{x})$ が全て $0$ で初期化されている。以下の種類のクエリが合計で $Q$ 個飛んでくるので、順に処理してください。

  • $\mathrm{add}(\bm{x}, v)$ : 格子点 $\bm{x}$ の重みに $v$ を加算する。即ち、$w(\bm{x})\leftarrow w(\bm{x}) + v$ とする。
  • $\mathrm{sum}(\bm{x}, \bm{y})$ : $\displaystyle \sum_{\bm{x} \leq \bm{z} \lt \bm{y}} w(\bm{z})$ を出力する。

ここで、$\bm{x}=(x_1,\ldots,x_D)$ および $\bm{y}=(y_1,\ldots,y_D)$ に対して、

$$\bm{x} \lt \bm{y} \defarrow \forall i \in \{1,\ldots,D\}.\; x_i\lt y_i$$

と定義します。その他の不等号 ($\leq,\geq,\gt$) に関しても同様に要素ごとの比較で定義します。


$D$ を定数とみなせば、この問題はクエリ先読みと $D$ 次元 Fenwick Tree により $O(Q (\log Q) ^ D)$ 時間で解くことができますが、一般の $\mathrm{sum}$ クエリに関しては定数倍として少なくとも $2 ^ D$ が掛かります (後述)。

なお、$\bm{x}=-\bm{\infty}\coloneqq (-\infty,\ldots,-\infty)$ であるような場合の $\mathrm{sum}$ クエリはこの $2 ^ D$ を落とすことができます。そのため、このケースを特別視して以降は $\mathrm{sum}(-\bm{\infty},\bm{y})$ を $\mathrm{prefix\_sum}(\bm{y})$ と書きます。

解法

$D=1$ の場合

この場合は、$\mathrm{add}$ クエリで現れる座標を先読みして座標圧縮を行うことで、通常の Fenwick Tree により全体 $O(Q\log Q)$ 時間で解くことができます。

$D\gt 1$ の場合

概要

$D=1$ の場合と同様に $\mathrm{add}$ クエリで現れる座標の $1$ 次元目を座標圧縮しておきます。そして、各ノードに $D-1$ 次元 Fenwick Tree を乗せることで $D$ 次元 Fenwick Tree を構築します。これにより、次で述べるような方法で各種クエリに答えることができます。

$3$ 次元 Fenwick Tree

上図は $D$ 次元 Fenwick Tree のノード $N_i$ が $D-1$ 次元 Fenwick Tree であり、即ち再帰的な構造を持っていることを示しています。

クエリ処理

一旦 $D$ 次元 Fenwick Tree の構築が完了したと思うことにして、$\mathrm{add}, \mathrm{sum}$ クエリの実装イメージを示します。通常のFenwick Treeを知っていれば難しくないと思います。

# 1-indexed
# n : ノード数
# N[i] := Fenwick Tree の i 番目のノード

def add((x1, ..., xD), v):
  index = compress(x1) # x1 を座標圧縮して得られる添字
  while index <= n:
    N[index].add((x2, ..., xD), v)
    index += index & -index

def sum((x1, ..., xD), (y1, ..., yD)):
  if x1 >= y1:
    return 0

  result = 0
  
  i = max_index_less_than(y1) # y1未満の最大の座標に対応する添字
  while i > 0:
    result += N[i].sum((x2, ..., xD), (y2, ..., yD))
    i -= i & -i

  i = max_index_less_than(x1) # x1未満の最大の座標に対応する添字
  while i > 0:
    result -= N[i].sum((x2, ..., xD), (y2, ..., yD))
    i -= i & -i

  return result

以降の計算量解析では $D$ は定数とし、また $\log$ の底は $2$ とします。

$\mathrm{add}$ クエリの時間計算量 $T _ {\mathrm{add}}(D)$ は

$$\begin{aligned} T _ {\mathrm{add}}(1) & = O(\log Q), \\ T _ {\mathrm{add}}(D) & \leq (\log Q + 1) \cdot (T _ {\mathrm{add}}(D-1) + O(1)) \quad (D \gt 1) \end{aligned}$$

より $T_{\mathrm{add}}(D) = O( (\log Q) ^ D)$ です *1

$\mathrm{prefix\_sum}$ クエリの時間計算量 $T _ {\mathrm{psum}}(D)$ は、sum において $2$ つ目の while を skip した場合の計算量なので

$$\begin{aligned} T _ {\mathrm{psum}}(1) &= O(\log Q), \\ T _ {\mathrm{psum}}(D) &\leq (\log Q + 1) \cdot (T _ {\mathrm{psum}}(D-1) + O(1)) \quad (D \gt 1) \end{aligned}$$

より $T _ {\mathrm{psum}}(D) = O( (\log Q) ^ D)$ です。

$\mathrm{sum}$ クエリの時間計算量 $T _ {\mathrm{sum}}(D)$ は

$$\begin{aligned} T _ {\mathrm{sum}}(1) &= O(\log Q), \\ T _ {\mathrm{sum}}(D) &\leq (2\log Q + 1) \cdot (T _ {\mathrm{sum}}(D-1) + O(1)) \quad (D \gt 1) \end{aligned}$$

より $T _ {\mathrm{sum}}(D) = O( (\log Q) ^ D)$ です。ただし、上式から分かるように、定数倍として少なくとも $2 ^ D$ が掛かります *2。実用する際は意識するとよいかもしれません。

構築

$\mathrm{add}$ クエリで与えられる格子点の列 $\bm{x} ^ 1,\ldots,\bm{x} ^ k$ をクエリ先読みにより得られているとします。

格子点 $\bm{x} ^ i = (x ^ i _ 1, \ldots, x ^ i _ D)$ の $1$ 次元目の値 $x ^ i _ 1$ を集めてできる (多重でない) 集合 (形式的には $\{ p \mid \exists i \in \{1, \ldots, k\} \; \text{s.t.}\; x ^ i _ 1 = p \}$) の要素を昇順に並べてできる列を $P=(p_1, \ldots, p_n)$ とすると、構築する $D$ 次元 Fenwick Tree のイメージ図は次のようになります。

格子点 $\bm{x} ^ i$ の情報を持つべきノード $N_j$ は $\mathrm{add}$ クエリを処理する場合と同様にして列挙できるので、以下のようにして構築パートを実装できます。

# 1-indexed

points = empty_list()

# クエリ先読みでFenwick Treeに点を登録する
def register_point((x1, ..., xD)):
  points.append((x1, ..., xD))

# p[i] は上で p_i として説明した変数
p = None

# 座標圧縮
def compress(x):
  l = 0
  r = p.size() + 1
  while r - l > 1:
    m = floor_div(l + r, 2)
    if p[m] < x:
      l = m
    else:
      r = m
  return r

def build():
  # pointsの1次元目の値のリスト
  coordinates = empty_list()
  for (x1, ..., xD) in points:
    coordinates.append(x1)
  # coordinatesの要素を昇順に、重複なしで並べたリスト
  p = sorted_unique_list(coordinates)
  n = p.size()

  # 再帰終端
  if D == 1:
    return

  for (x1, ..., xD) in points:
    # p[index]=x1 なるindex
    i = compress(x1)
    # addクエリと同様にして i の親を辿っていく
    while i <= n:
      # N[i]は(D-1)次元Fenwick Tree
      N[i].register_point((x2, ..., xD))
      i += i & -i

  for i in {1, ..., n}:
    N[i].build()

構築にかかる時間計算量を考えます。各点に対して、最も内側の while が回る回数は $O( (\log Q) ^ {D - 1} )$ 回です。従って、$p$ の計算にかかる時間計算量はソートがボトルネックとなり

$$ O( (Q (\log Q) ^ {D - 1} ) \log (Q (\log Q) ^ {D - 1} )) = O(Q (\log Q) ^ {D - 1} (\log Q + D \log \log Q)) $$

と評価できます。$D$ は定数なので、これは $O(Q (\log Q) ^ D)$ です。また、座標圧縮を行う部分の時間計算量は $O(Q \cdot ( (\log Q) ^ {D - 2} \log Q)) = O(Q (\log Q) ^ {D - 1})$ です。

以上より、構築にかかる時間計算量は $O(Q (\log Q) ^ D)$ です。while が回る回数から、空間計算量は $O(Q (\log Q) ^ {D - 1})$ と評価されます。

実装

ここまで python 風の擬似コードを書いておいてなんですが、C++ による実装例を以下に示します。

折り畳み

  • 動作は保証しません
  • 説明とは異なり全て0-indexedです
#include <algorithm>
#include <array>
#include <cassert>
#include <limits>
#include <vector>

namespace compressed_fenwick_tree {
    namespace detail {
        // 座標圧縮
        template <typename CoordinateType>
        struct Compressor {
            using coordinate_type = CoordinateType;

            Compressor() = default;

            void register_coordinate(const coordinate_type& x) {
                xs.push_back(x);
            }

            int build() {
                std::sort(xs.begin(), xs.end());
                xs.erase(std::unique(xs.begin(), xs.end()), xs.end());
                return xs.size();
            }

            // x以上の最小の値の添字(存在しない場合はxs.size())
            int operator()(const coordinate_type& x) const {
                return std::lower_bound(xs.begin(), xs.end(), x) - xs.begin();
            }
        private:
            std::vector<coordinate_type> xs;
        };
    }

    /**
     * @brief D>1の場合のD次元Fenwick Tree
     *
     * @param T 要素型
     * @param op 二項演算の結果を返す関数
     * @param e 単位元を返す関数
     * @param inv 逆元を返す関数
     * @param D 次元
     * @param Index 座標の型
     */
    template <typename T, T(*op)(T, T), T(*e)(), T(*inv)(T), std::size_t D = 1, typename Index = int>
    struct CompressedFenwickTree {
        using value_type = T;
        using index_type = Index;
        // 点を表す型
        using point_type = std::array<index_type, D>;
        // D次元の直方体領域を表す型
        using cube_type = std::array<std::pair<index_type, index_type>, D>;
        // ノードの型
        using node_type = CompressedFenwickTree<value_type, op, e, inv, D - 1, index_type>;

        CompressedFenwickTree() = default;

        /**
         * @brief addクエリで渡される点を予め与えておく
         *
         * @param p 登録する点
         */
        void register_point(const point_type& p) {
            compressor.register_coordinate(p[0]);
            points.push_back(p);
        }

        /**
         * @brief Fenwick Treeを構築する (クエリを投げる前に必ず呼ぶ)
         */
        void build() {
            nodes.assign(n = compressor.build(), node_type{});
            for (const auto& p : points) for (int k = compressor(p[0]) + 1; k <= n; k += k & -k) {
                nodes[k - 1].register_point(tail(p));
            }
            points.clear();
            points.shrink_to_fit();
            for (auto& t : nodes) t.build();
        }

        /**
         * @brief D次元直方体領域に含まれる点の重みの総和を計算する
         *
         * @param cube D次元直方体領域
         */
        value_type sum(const cube_type& cube) const {
            return op(sum(cube[0].second, tail(cube)), inv(sum(cube[0].first, tail(cube))));
        }

        /**
         * @brief 点の重みに加算する
         *
         * @param p 重みを加算する点 (必ずregister_pointで予め渡した点を与える。それ以外の点に対する動作は未定義)
         * @param val 加算する値
         */
        void add(const point_type& p, const value_type& val) {
            for (int r = compressor(p[0]) + 1; r <= n; r += r & -r) {
                nodes[r - 1].add(tail(p), val);
            }
        }
    private:
        // ノードの数
        int n;
        // 1次元目の値を座標圧縮するための構造体
        detail::Compressor<index_type> compressor;
        // addクエリが来る点のリスト
        std::vector<point_type> points;
        // ノード
        std::vector<node_type> nodes;

        /**
         * @brief D次元直方体領域に含まれる点の重みの総和を計算する。(1次元目に下限が無い場合)
         *
         * @param head_r 1次元目の上限
         * @param p 2次元目以降を表すD-1次元直方体領域
         */
        value_type sum(const index_type& head_r, const typename node_type::cube_type& p) const {
            if (head_r == std::numeric_limits<index_type>::min()) return e();
            value_type res = e();
            for (int r = compressor(head_r); r; r -= r & -r) {
                res = op(res, nodes[r - 1].sum(p));
            }
            return res;
        }

        /**
         * @brief array<point_type>またはarray<cube_type>の第二要素以降を取り出す
         *
         * @param p 点またはD次元直方体領域
         */
        template <typename X>
        static constexpr auto tail(const X& p) {
            return tail_impl(p, std::make_index_sequence<D - 1>{});
        }
        template <typename X, std::size_t... Seq>
        static constexpr auto tail_impl(const X& p, std::index_sequence<Seq...>) {
            using return_type = std::conditional_t<
                std::is_same_v<X, point_type>,
                typename node_type::point_type, // Xがpoint_typeの場合
                typename node_type::cube_type   // Xがcube_typeの場合
            >;
            return return_type{ p[1 + Seq]... };
        }
    };

    /**
     * @brief 1次元Fenwick Tree
     *
     * @param T 要素型
     * @param op 二項演算の結果を返す関数
     * @param e 単位元を返す関数
     * @param inv 逆元を返す関数
     * @param Index 座標の型
     */
    template <typename T, T(*op)(T, T), T(*e)(), T(*inv)(T), typename Index>
    struct CompressedFenwickTree<T, op, e, inv, std::size_t(1), Index> {
        using value_type = T;
        using index_type = Index;
        using point_type = index_type;
        using cube_type = std::pair<index_type, index_type>;
        using node_type = value_type;

        CompressedFenwickTree() = default;

        /**
         * @brief addクエリで渡される点を予め与えておく
         *
         * @param p 登録する点
         */
        void register_point(const point_type& p) {
            compressor.register_coordinate(p);
        }

        /**
         * @brief Fenwick Treeを構築する (クエリを投げる前に必ず呼ぶ)
         */
        void build() {
            nodes.assign(n = compressor.build(), e());
        }

        /**
         * @brief 半開区間内の重みの総和を計算する
         *
         * @param l 区間の左端 (閉)
         * @param r 区間の右端 (開)
         */
        value_type sum(const index_type& l, const index_type& r) const {
            return op(sum(r), inv(sum(l)));
        }

        /**
         * @brief 半開区間内の重みの総和を計算する
         *
         * @param range 半開区間
         */
        value_type sum(const cube_type& range) const {
            return sum(range.first, range.second);
        }

        /**
         * @brief 点の重みに加算する
         *
         * @param p 重みを加算する点 (必ずregister_pointで予め渡した点を与える。それ以外の点に対する動作は未定義)
         * @param val 加算する値
         */
        void add(const point_type& p, const value_type& val) {
            for (int r = compressor(p) + 1; r <= n; r += r & -r) {
                nodes[r - 1] = op(nodes[r - 1], val);
            }
        }
    private:
        int n;
        detail::Compressor<index_type> compressor;
        std::vector<node_type> nodes;

        /**
         * @brief prefix query
         *
         * @param p 区間の右端 (開)
         */
        value_type sum(const point_type& p) const {
            if (p == std::numeric_limits<index_type>::min()) return e();
            value_type res = e();
            for (int r = compressor(p); r; r -= r & -r) {
                res = op(res, nodes[r - 1]);
            }
            return res;
        }
    };

    namespace detail {
        template <typename T, T(*e)()>
        T inv_dummy(T) { return e(); }
    }

    /**
     * @brief 任意次元Fenwick Treeで、sumクエリがprefix sumのみの場合
     *
     * @param T 要素型
     * @param op 二項演算の結果を返す関数
     * @param e 単位元を返す関数
     * @param D 次元
     * @param Index 座標の型
     */
    template <typename T, T(*op)(T, T), T(*e)(), std::size_t D = 1, typename Index = int>
    struct CompressedFenwickTreePrefixQuery : private CompressedFenwickTree<T, op, e, detail::inv_dummy<T, e>, D, Index> {
        using base_type = CompressedFenwickTree<T, op, e, detail::inv_dummy<T, e>, D, Index>;

        using value_type = typename base_type::value_type;
        using index_type = typename base_type::index_type;
        using point_type = typename base_type::point_type;

        using base_type::base_type;
        using base_type::register_point;
        using base_type::build;
        using base_type::add;

        /**
         * @brief prefix sumを計算する
         *
         * @param p 領域の右端
         */
        value_type sum(const point_type& p) {
            return base_type::sum(make_cube(p));
        }

    private:
        using cube_type = typename base_type::cube_type;

        cube_type make_cube(const point_type& p) {
            static constexpr index_type neg_inf = std::numeric_limits<index_type>::min();
            if constexpr (D == 1) {
                return cube_type{ neg_inf, p };
            } else {
                cube_type cube;
                for (std::size_t i = 0; i < D; ++i) cube[i] = { neg_inf, p[i] };
                return cube;
            }
        }
    };
} // namespace compressed_fenwick_tree

使用例

結び

多分どこか間違えてるので指摘してください。

*1: $T _ {\mathrm{add}}(D) \leq \sum _ {i = 0} ^ {D} c ^ D _ i (\log Q) ^ i$ なる係数列 $c ^ D$ を $Q$ の値に依存しないように $D$ が小さい順に決めることができて、これに対して十分大きな係数 $C _ D$ を取って $\forall x \gt 0.\; \sum _ {i = 0} ^ {D} c ^ D _ i x ^ i \leq C _ D x ^ D$ を満たすようにできます。これは当然 $D$ に (のみ) 依存しますが、今回の解析では $D$ は定数なので $C_D$ も定数です。従って、$D$ を定数とすれば $T _ {\mathrm{add} }(D) = O( (\log Q) ^ D)$ と結論付けられます。そして、これは $T _ {\mathrm{sum} }(D)$ や $T _ {\mathrm{prefix\_ sum} }(D)$ に関しても同様です。

*2:この $2 ^ D$ は $\mathrm{prefix\_sum}$ を用いて包除原理で超直方体領域内の和を計算していることで掛かる係数と解釈できます。