√w-bit の入力に対する定数時間 rank/select

これいる？

背景
本題
実測
雑記
あとがき
おわり

背景

定数時間 rank/select をできる簡潔データ構造を作るときの基本戦略として、表引き (table lookup) というのがあります。小さいサイズ（たとえば $\log_2(n)/2$-bit 程度）の入力での答えを予め計算して配列に入れておき、必要に応じてそれを返すというものです。 rank/select のように引数が一個の問題であれば、$2^{\log_2(n)/2}\cdot \log_2(n)/2$ 個のエントリの配列が作られます。各答えは $\log_2(\log_2(n)/2)$ bits で表せるので、 $$ \begin{aligned} 2^{\log_2(n)/2}\cdot \log_2(n)/2 \cdot \log_2(\log_2(n)/2) &\le \sqrt{n}\cdot \log_2(n) \cdot \log_2(\log_2(n)) \\ &= o(n) \end{aligned} $$ となり、$o(n)$ bits で済むので大丈夫という感じです。

とはいえ、毎回毎回配列にアクセスするのは嫌じゃない？という気持ちがあります。

本題

ということで、別の手法を考えます。

俺はたった今から配列を捨てる！

配列を使わないとなったら、ビット演算などでがちゃがちゃするくらいしかないので、そうします。

rsk0315.hatenablog.com

例としては $w = 64$, $\sqrt w=8$ のものを挙げますが、ワードサイズが定数という意味ではないので気をつけましょう。すなわち、$w$ 回のループを回す $\Theta(w)$ 時間のアルゴリズムや、$w$ に対して二分探索する $\Theta(\log(w))$ 時間のアルゴリズムは定数時間とは認められません。

ビット並列

$\gdef\code#1{\text{\texttt{#1}}}$ $\gdef\hexn#1{{\code{#1}}_{(16)}}$ $\gdef\binn#1{{\code{#1}}_{(2)}}$

まず、$w$-bit の word に対して、「$\sqrt w$-bit からなる block が $\sqrt w$ 個ある」という見方をします。たとえば、$\hexn{B68D2D79D3D9821A}$ という値は $$\angled{\hexn{1A}, \hexn{82}, \hexn{D9}, \hexn{D3}, \hexn{79}, \hexn{2D}, \hexn{8D}, \hexn{B6}}$$ という block 列と見做します（下位ビット側を先頭側とします）。

サブルーチン

さて、いくつかの基本的な演算を導入します。

$\sqrt w$-bit の $u$ に対して、各 block が $u$ であるような word を返す演算を splat とします*1。たとえば $$\code{splat}(\hexn{A3}) = \angled{\hexn{A3}, \hexn{A3}, \hexn{A3}, \hexn{A3}, \hexn{A3}, \hexn{A3}, \hexn{A3}, \hexn{A3}}$$ です。

block の列を表す word $w$ に対して、$w$ 中の対応する block が $0$ なら $0$、そうでないなら $1$ となるような word を返す演算を nonzero とします。たとえば $$ \begin{aligned} &\phantom{{}={}} \code{nonzero}(\angled{\hexn{B3}, \hexn{6C}, \hexn{00}, \hexn{53}, \hexn{C6}, \hexn{00}, \hexn{F1}, \hexn{5D}}) \\ &= \angled{\hexn{01}, \hexn{01}, \hexn{00}, \hexn{01}, \hexn{01}, \hexn{00}, \hexn{01}, \hexn{01}} \end{aligned} $$ です。

$\sqrt w$-bit の $u$ に対して、$i$ 番目の block が「$u$ の $i$ 番目の bit が $0$ のとき $0$、そうでないとき $1$」となるような word を返す演算を expand とします。たとえば $$ \begin{aligned} &\phantom{{}={}} \code{expand}(\hexn{B5}) \\ &= \code{expand}(\binn{10110101}) \\ &= \angled{\hexn{01}, \hexn{00}, \hexn{01}, \hexn{00}, \hexn{01}, \hexn{01}, \hexn{00}, \hexn{01}} \end{aligned} $$ です（下位 bit が block の先頭側であることに注意）。

block 列を表す word $w$ であって、各 block が $0$ または $1$ であるものに対して、$i$ 番目の block が「$i$ 番目以下の block の $1$ の個数」となるような word を返す演算を accumulate とします。たとえば $$ \begin{aligned} &\phantom{{}={}} \code{accumulate}(\angled{\hexn{01}, \hexn{00}, \hexn{01}, \hexn{00}, \hexn{01}, \hexn{01}, \hexn{00}, \hexn{01}}) \\ &= \angled{\hexn{01}, \hexn{01}, \hexn{02}, \hexn{02}, \hexn{03}, \hexn{04}, \hexn{04}, \hexn{05}} \end{aligned} $$ です。

block 列を表す word $w$ と、$\sqrt w$ 未満の非負整数 $i$ に対して、$i$ 番目の block の値を返す演算を get とします。たとえば $$ \begin{aligned} &\phantom{{}={}} \code{get}(\angled{\hexn{B3}, \hexn{6C}, \hexn{00}, \hexn{53}, \hexn{C6}, \hexn{00}, \hexn{F1}, \hexn{5D}}, 6) \\ &= \hexn{F1} \end{aligned} $$ です。

block 列を表す word $w_L$, $w_R$ であって、各 block が $2^{\sqrt w-1}$ 未満であるものに対して、$i$ 番目の block が「$\code{get}(w_L, i) \ge \code{get}(w_R, i)$ なら $1$、そうでないなら $0$」となるような word を返す演算を gt-eq とします。たとえば、 $$ \begin{aligned} &\code{gt-eq}( \\ & \begin{aligned} &&&&& \angled{\hexn{41}, \hexn{31}, \hexn{43}, \hexn{65}, \hexn{31}, \hexn{10}, \hexn{61}, \hexn{67}}, \\ &&&&& \angled{\hexn{01}, \hexn{36}, \hexn{44}, \hexn{15}, \hexn{44}, \hexn{77}, \hexn{61}, \hexn{27}} \end{aligned} \\ &)= \angled{\hexn{01}, \hexn{00}, \hexn{00}, \hexn{01}, \hexn{00}, \hexn{00}, \hexn{01}, \hexn{01}} \end{aligned} $$ です。

block 列を表す word $w$ に対して、$i$ 番目の block が「$i = 0$ なら $0$、そうでなければ $\code{get}(w, i-1)$」となるような演算を shift とします。たとえば $$ \begin{aligned} &\phantom{{}={}} \code{shift}(\angled{\hexn{41}, \hexn{31}, \hexn{43}, \hexn{65}, \hexn{31}, \hexn{10}, \hexn{61}, \hexn{67}}) \\ &= \angled{\hexn{00}, \hexn{41}, \hexn{31}, \hexn{43}, \hexn{65}, \hexn{31}, \hexn{10}, \hexn{61}} \end{aligned} $$ です。

block 列を表す word $w$ であって、各 block が $0$ または $1$ であるものに対して、$1$ であるような block の個数を返す演算を popcnt とします。たとえば $$ \begin{aligned} \code{popcnt}(\angled{\hexn{01}, \hexn{00}, \hexn{00}, \hexn{01}, \hexn{00}, \hexn{00}, \hexn{01}, \hexn{01}}) = 4 \end{aligned} $$ です。

メイン

これらによって、$\code{rank}(u, i)$ と $\code{select}(u, i)$ は次のように書けます。ただし、$\code{rank}(u, i)$ は $u$ のうち $i$ 番目未満の bit における $1$ の個数、$\code{select}(u, i)$ は $u$ のうち $i$ 番目 (0-indexed) の $1$ の位置を表します。

$$ \begin{aligned} \code{rank}(u, i) &= \code{get}( (\code{shift}\circ\code{accumulate}\circ\code{expand})(u), i), \\ \code{select}(u, i) &= \code{popcnt}(\code{gt-eq}(\code{splat}(i), (\code{accumulate}\circ\code{expand})(u))). \end{aligned} $$

実装

サブルーチンの実装について述べます。まず定数を用意します。

$\gdef\Mlow{M_{\text{lo}}}$ $\gdef\Mhigh{M_{\text{hi}}}$ $\gdef\Miota{M_{\text{iota}}}$

$$ \begin{aligned} \Mlow &= \frac{2^w-1}{2^{\sqrt w}-1} \\ &= \angled{0, 0, \dots, 0}, \\ \Mhigh &= 2^{\sqrt w-1}\times \Mlow \\ &= \angled{\sqrt w-1, \sqrt w-1, \dots, \sqrt w-1}, \\ \Miota &= \frac{2^{w+\sqrt w}-1}{2^{\sqrt w+1}-1} \\ &= \angled{0, 1, \dots, \sqrt w-1}. \end{aligned} $$

$\gdef\bitand{\wedge}$ $\gdef\bitor{\vee}$ $\gdef\bitnot{\lnot}$

これを用いて、次のように表せます。 $$ \begin{aligned} \code{splat}(u) &= \Mlow\times u, \\ \code{nonzero}(w) &= \Floor{\frac{(w\bitor \bitnot(\Mhigh - (w \bitand \bitnot \Mhigh))) \bitand \Mhigh}{2^{\sqrt w-1}}} \\ &= \Floor{\frac{(w\bitor (\Mhigh - \Mlow + w)) \bitand \Mhigh}{2^{\sqrt w-1}}}, \\ \code{expand}(u) &= \code{nonzero}(\code{splat}(u) \bitand \Miota), \\ \code{accumulate}(w) &= (w \times \Mlow)\bmod 2^w, \\ \code{get}(w, i) &= \Floor{\frac w{2^{i\sqrt w}}} \bitand (2^{\sqrt w}-1), \\ \code{gt-eq}(w_L, w_R) &= \Floor{\frac{( (w_L\bitor \Mhigh) - w_R) \bitand \Mhigh}{2^{\sqrt w-1}}}, \\ \code{shift}(w) &= (w\times 2^{\sqrt w})\bmod 2^w, \\ \code{popcnt}(w) &= \Floor{\frac{\code{accumulate}(w)}{2^{w-\sqrt w}}}. \end{aligned} $$

ただし、$\bitand$, $\bitor$, $\bitnot$ はそれぞれ bitwise-AND (&), bitwise-OR (|), bitwise-NOT (!, ~) を表します。また、$(x \times y)\bmod 2^w$ はオーバーフローを無視する乗算 (wrapping_mul)、$\floor{x/2^i}$ や $(x\times 2^i)\bmod 2^w$ はシフト演算 (>>, <<) に相当します。

gt-eq と nonzero が難しいです。$\sqrt w-1$ 番目の bit とそれ以外で分けて考えたり、分配法則を考えつつ演算を省略したりしています。$\sqrt w-1$ 番目の bit を、減算の繰り下がり（あるいは加算の繰り上がり）に関する bit として使っています。

（追記）ここバグっているかも。引数の $\sqrt w-1$ 番目の bit が $1$ なのを仮定している？

また、$2^i-1$ は (1 << i) - 1 ではなく !(!0 << i) とも書けます。括弧を端折れがちなのでうれしい場合があります。

実装例 (Rust)

const W: usize = u64::BITS as usize;
const BLOCK_LEN: usize = 8;
const BLOCK: u64 = !(!0 << BLOCK_LEN);
const HI_POS: usize = BLOCK_LEN - 1;
const M_LO: u64 = 0x0101010101010101;
const M_IOTA: u64 = 0x8040201008040201;
const M_HI: u64 = M_LO << HI_POS;

#[inline(always)]
fn splat(w: u8) -> u64 { M_LO * w as u64 }
#[inline(always)]
fn nonzero(w: u64) -> u64 { ((w | (M_HI - M_LO + w)) & M_HI) >> HI_POS }
// fn nonzero(w: u64) -> u64 { ((w | !(M_HI - (w & !M_HI))) & M_HI) >> HI_POS }
#[inline(always)]
fn expand(w: u8) -> u64 { nonzero(splat(w) & M_IOTA) }
#[inline(always)]
fn accumulate(w: u64) -> u64 { w.wrapping_mul(M_LO) }
#[inline(always)]
fn get(w: u64, i: usize) -> usize { (w >> (BLOCK_LEN * i) & BLOCK) as _ }
#[inline(always)]
fn gt_eq(wl: u64, wr: u64) -> u64 { (((wl | M_HI) - wr) & M_HI) >> HI_POS }
#[inline(always)]
fn shift(w: u64) -> u64 { w << BLOCK_LEN }
#[inline(always)]
fn popcnt(w: u64) -> usize { (accumulate(w) >> (W - BLOCK_LEN)) as _ }

#[inline(always)]
pub fn rank(w: u8, i: usize) -> usize { get(shift(accumulate(expand(w))), i) }
#[inline(always)]
pub fn select(w: u8, i: usize) -> usize {
    popcnt(gt_eq(splat(i as _), accumulate(expand(w))))
}

テスト

#[cfg(test)]
mod tests {
    use crate::*;

    const fn rank_table<const LEN: usize, const PAT: usize>() -> [[u8; LEN]; PAT]
    {
        let mut res = [[0; LEN]; PAT];
        let mut i = 0;
        while i < PAT {
            let mut cur = 0;
            let mut j = 0;
            while j < LEN {
                res[i][j] = cur;
                if i >> j & 1 != 0 {
                    cur += 1;
                }
                j += 1;
            }
            i += 1;
        }
        res
    }

    const fn select_table<const LEN: usize, const PAT: usize>()
    -> [[u8; LEN]; PAT] {
        let mut res = [[0; LEN]; PAT];
        let mut i = 0;
        while i < PAT {
            let mut cur = 0;
            let mut j = 0;
            while j < LEN {
                if i >> j & 1 != 0 {
                    res[i][cur] = j as _;
                    cur += 1;
                }
                j += 1;
            }
            i += 1;
        }
        res
    }

    const RANK_TABLE: [[u8; 8]; 256] = rank_table::<8, 256>();
    const SELECT_TABLE: [[u8; 8]; 256] = select_table::<8, 256>();

    #[test]
    fn test_rank() {
        for w in 0_u8..=!0 {
            for i in 0..8 {
                assert_eq!(rank(w, i), RANK_TABLE[w as usize][i] as usize);
            }
        }
    }

    #[test]
    fn test_select() {
        for w in 0_u8..=!0 {
            for i in 0..w.count_ones() as _ {
                assert_eq!(select(w, i), SELECT_TABLE[w as usize][i] as usize);
            }
        }
    }
}

サイズに関して

よく見かけるものでは、rank の小ブロックは $\log_2(n)$ bits に設定されがちですが、必ずしもそうする必要はありません。大ブロックを $l = \omega(\log(n))$ bits として $s = \omega(\log(l))$ bits であればよいはずです。

$s \le \sqrt w$ とするためには $w \ge s^2$ が必要です。ワードサイズの都合から $w\ge \log_2(n)$ なので、$\log_2(n)\ge s^2$ なら十分です。あるある設定に従って $l = \log_2(n)^2$ とすると、$\log_2(n)^{1/2} = \omega(\log(\log_2(n)^2))$ なので、$s = \log_2(n)^{1/2}$ とすればよさそうです。もちろん $l$ の選び方にもよる（よくある $\log_2(n)^2$ にこだわる必要はない）ので、その辺はよしなにという感じです。

select の小ブロックについても別途考える必要がありますが、解析自体は同じようにできて、$s = \log_2(n)^{1/2}$ で大丈夫そうです。

実測

やったので追記です。

rank は $2^8\cdot 8$ 通り、select は $\tfrac12\cdot 2^8\cdot8$ 通り、有効なクエリがあります。有効なクエリを予め生成して、rank は 100 周、select は 200 周ぶんだけランダムに与えました。

計測の対象は下記の 3 つです。

コンパイル時計算で [[u8; 8]; 256] に答えを入れたもの (const-table)
実行時計算で IntVec に答えを入れたもの (runtime-table)
今回のもの (calculate)

IntVec は、（各々の答えを持っておくのに 8-bit 整数を使うのは贅沢なので）必要な分の bit 数ずつ保持している自前の struct です。取り出す際に適切な word を高々 2 つ持ってくる必要があるので比較的オーバーヘッドが大きいです。

対象	rank	select
const-table	97.572 µs	93.543 µs
runtime-table	350.70 µs	352.93 µs
calculate	134.72 µs	180.47 µs

まぁ悪くないのではないでしょうか（甘めの評価）。

雑記

ここ追記です。

基本的なサブルーチンいくつかで書けることを示したので、そういう SIMD の命令が使える文脈ではうまく高速化できる場合もあるかもしれません。オーバーヘッドの方が大きいかも。

また、算術演算・ビット演算で書けることを示したので、そういう RAM（word RAM ではなく任意長の整数を扱えるもの）において bitvector を一つの整数として持つことで $O(1)$ time, $O(1)$ extra space で処理できてしまうことがわかります。これは RAM の abuse であって、word RAM で考えることが妥当だと主張するのに役立ちそうな気がします。

あとがき

おふとんでごろごろしているときに考えたので、あまりサーベイをしていません。たぶん人々にいろいろ考えられていそう。あと実測をまだしていないのですが、表引きの方が結果的に速そうな感じがしませんか。かなしかなし。

（追記）→ まぁよし？

とはいえ、「$\sqrt w$-bit の入力に対してがちゃがちゃやる」系のパズルを考えるの自体はたのしいのでよいです。「空間が $o(n)$ bits に収まるようにがちゃがちゃやる」系のパズルを考えるのがたのしいのと似ています。

実測が速いとうれしい気持ちにはなるものの、実測にはそこまで興味がないので、結果的にこういうことばかり考えてしまいます。

途中でリンクを貼った記事では $\sqrt{w}$-bit の入力に対する msb を求めましたが、msb は select の特殊ケースなんですよね。$u$ の bit のうち $1$ のもの (popcount) を $k$ として $\code{msb}(u) = \code{select}(u, k-1)$ なので。また、この $k$ も $(\code{popcnt}\circ\code{expand})(u)$ などで求められます。個数を数えたいだけであれば、$\code{nonzero}$ を経由せずに $\code{splat}(u)\bitand \Miota$ に適当な定数を掛けてシフトするのでも十分そうです。

おわり

日曜日返して

*1:SIMD の文脈でそういう言い方をするらしいので。

えびちゃんの日記

えびちゃん（競プロ）の日記です。

√w-bit の入力に対する定数時間 rank/select

背景

本題

ビット並列

サブルーチン

メイン

実装

サイズに関して

実測

雑記

あとがき

おわり