Wavelet Matrix 3D

3 次元クエリに対応できる Wavelet Matrix を実装する。

目的と実装方法

目的

3 次元での $k$-th クエリの問題($x_0 \le x < x_1$, $y_0 \le y < y_1$ の範囲に存在するデータについて、$k$ 番目に小さい $z$ の値を求める問題)に、なるべく省メモリかつ高速に答えられるようにする。

(静的 ?)3 次元クエリは 永続データ構造 と Wavelet Matrix を組み合わせて解けることが知られている:

今回は永続データ構造を使わずに実装する。

[追記]

実装方法

2 次元範囲のデータの個数を取得できる Wavelet Matrix を各ノードとする 2 分木を作ると、3 次元での $k$-th クエリに対応できる。Fenwick Tree での 2 分探索のような形で探索するため、$z$ 軸の各 bit の階層において、片側の bit のノードのみを配置すれば大丈夫。

入力が、サイズ $X$ の $(y, z)$ ($0 \le y < Y$, $0 \le z < Z$) の配列であったとき、

  • 領域: $O(X \log{Y} \log{Z})$ : ほぼ $(4 \cdot X \log{Y} \log{Z})$ Bytes。
  • 構築: $O(X \log{Y} \log{Z})$
  • クエリ: $O(\log{Y} \log{Z})$

になる。

上の実装は non-succinct な Wavelet Matrix を使った場合で、compact (succinct) な方を用いれば $O(X \log{Z} + f(X, Z) \log{Y})$ 領域で収まる。ここで、$f(X, Z)$ は 2 分木のノードの個数で、

  • $X \ge Z$ のとき、$f(X, Z) \in O(Z)$
  • $X < Z$ のとき、$f(X, Z) \in O((\log(Z/X) + 1) X)$

となる。$z$ 軸の要素を座標圧縮できるならば、$\log$ は(実質) 1 つしか乗らない。

ソースコード

  • 実装 1: $O(X \log{Y} \log{Z})$ space

構築中に malloc と free を発生させないようにする必要があり、これが恐らく一番面倒。

[2017-09-14: 追記]
Wavelet Matrix とかいているけど、この実装は Wavelet Tree に相当するようだ。

  • 実装 2: $O(X \log{Z} + f(X, Z) \log{Y})$ space

内部クラスの方は Wavelet Tree から Wavelet Matrix に変更した。クエリ関連の命令数が減ったため早くなると思ったのだけど、$N$ が小さいときは早くなるが、$N$ が大きいとき($N$ > 500000 など)は(実験例に対して)1 割ほど遅くなった。

外側のクラスの方は Wavelet Tree の構造のままだけど、2 分木を作りやすいため、そのままにしている。こちらも compact となる実装に変更し、メモリ使用量を減らした。

実行例

$A$ の要素の範囲を $[1, N/5]$、$a$, $b$, $c$, $d$ の範囲を $[0, N]$、$k$ の範囲を $[1, N/15]$ にした例。

$N$ $Q$ 構築(秒) クエリ(秒)
100000 100000 0.122 0.490
500000 500000 0.739 7.09
1000000 1000000 1.66 17.7
5000000 5000000 10.6 134.4

構築時間をこれ以上改善するのは難しいと思う。

$x$ 座標でソートした場合の解法。

改善箇所

  • Wavelet Tree の各ノードは、$z$ の座標の bit の $\{0, 1\}$で篩い分けた際に、個数が少なかった方のみを採用できる(今は bit が 0 の方を残している)。そのように変更すれば、少し早くなるかもしれない。
  • もしくは、Wavelet Tree の各ノードに境界値を陽に持たせて、探索が最適になるようにデータ配分してもいいのかもしれない。Wavelet Tree の定義に合致しなくなってしまうけど、メモリの節約法は引き継げる。