この記事は夏のブログリレー12日目の記事です

概要

bitset を使ったワードサイズ倍の高速化は、ほかの方法では計算量を改善しにくい問題で有用です。$N\times N$ サイズの $0$-$1$ 行列 $A$ をサイズ $N$ の bitset $N$ 個により保持しているとしましょう。$j$ を固定して $A_{i,j}$ に関する何らかの情報を得るためには、自明な実装ではワードサイズ倍の高速化を行えません。
しかし、転置を行うことにより単一のビット列の場合に帰着して、高速化を行うことができるようになります。この記事では、転置を高速に行うアルゴリズムおよび、転置を必要とせず、より高速に行える場合を紹介します。

以下では簡単のため正方行列のみを扱いますが、そうでない行列にも容易に適用可能です。

記法

添字は $0$-indexed とします。
$N\times N$ の $0$-$1$ 行列 $A$ の第 $i$ 行を $A_i$ とし、$N$ ビットのビット列として扱います。

ビット列 $x$ の第 $j$ ビットを $x_j$ と書きます。ビット列に対する $\wedge ,\vee ,\oplus ,\neg$ は、それぞれ各ビットごとの AND、OR、XOR、NOT を表します。
word size を $w$ とおいて $w=\Theta (\log N)$ と仮定します。

説明を簡単にするため $N$ を $2$ 冪と仮定し、$l={\log }_{2}N$ とします。
多重集合の和を $\uplus$ で表します。

目次

• 1. 高速な転置
• 2. 転置なしの高速化
  • 2.1 find_first
  • 2.2 popcount
• 3. 適用例

1. 高速な転置

問題

$N\times N$ の $0$-$1$ 行列 $A$ が与えられる。
各行は長さ $N$ のビット列として保持されているものとする。
この行列の転置 $A^{\mathsf{T}}$ を求めよ。

解法

転置では、$(r,c)$ にあるビットを $(c,r)$ に移せばよいです。
行番号と列番号を $2$ 進数で見ると、これはそれぞれのビットを入れ替えることに対応します。

次のアルゴリズムが正当です。

• $i=l-1,\dots ,0$ について次を行う。

  • 整数の組 $(r,c)$ であって、$r$ の $i$ ビット目が $0$ であり、$c$ の $i$ ビット目が $1$ であるもの全てに対して次を行う。

    • $(r,c)$ と $(r+2^i,c-2^i)$ にあるビットを交換する。

アルゴリズムの正当性

最初、$(r,c)$ にあったビットの動きを観察しましょう。
$r,c$ を $2$ 進数に展開して、$r_{l-1}{r}_{l-2}\dots r_0$ および $c_{l-1}{c}_{l-2}\dots c_0$ とします。

$i=k$ の場合、ビットは次のように動きます。

[1] $r_k=c_k$ のとき

変わりません。

[2] $r_k=0,c_k=1$ のとき

このとき、$r_k=1,c_k=0$ になります。また、$r_k=0,c_k=1$ なので、$r+2^k,c-2^k$ では繰り上げ、繰り下げが起こらず、$r_k,c_k$ 以外は変わりません。

[3] $r_k=1,c_k=0$ のとき

[2] で述べた交換のもう一方の要素となるため、$r_k=0,c_k=1$ になります。

以上から最終的に $r_i,c_i$ は全て swap されて転置されます。

具体例

$N=8$ の場合の転置の過程を示しておきます。

ビット演算による高速化

さらに、ビット演算により各 $i,r$ に対しての交換をまとめて行うことができます。
具体的には、次のような疑似コードで記述することができます。

for i = l - 1 to 0
    s = 1 << i

    # m_j = 1 ⇔ j の i 番目のビットが 0
    m = ...

    for r = 0 to N - 1
        if (r >> i) & 1:
            continue

        a = bs[r]
        b = bs[r + s]

        t = ((a >> s) ^ b) & m

        a ^= t << s
        b ^= t

        bs[r] = a
        bs[r + s] = b

特に、

t = ((a >> s) ^ b) & m

a ^= t << s
b ^= t

の部分の正当性について確認しておきます。

$j$ の $i$ ビット目が $0$ であるとします。この時、交換すべきビットの組は $a_{j+s}$ と $b_j$ です。
ここで、$a$ を $s$ だけ右シフトしていることに注意して、$t_j=a_{j+s}\oplus b_j$ となります。($\&m$ はコード中の $m$ の定義から干渉しないことに注意してください)
$x=a_{j+s},y=b_j$ とすると、次のように更新されており、$a_{j+s}$ と $b_j$ が swap されていることが確認できます。

$a_{j+s}\to x\oplus (x\oplus y)=y$

$b_j\to y\oplus (x\oplus y)=x$

交換する必要のないビットは $\&m$ により $t$ の対応するビットが $0$ となるため、値が変化しません。
以上からこのアルゴリズムは正当です。

計算量解析

ループは全体で $O(Nl)$ 回実行され、ループ内で行われる操作はいずれも長さ $N$ のビット列に対する演算であるため、1 回あたり $O\left(\frac{N}{w}\right)$ 時間です。
以上から、このアルゴリズムの計算量は $O\left(\frac{N^{2}l}{w}\right)$ です。

詳しくは記述しませんが、$w\times w$ ごとに分割して処理することで $O\left(\frac{N^2\log w}{w}\right)$ 時間にすることもできます。

2. 転置なしの高速化

転置を行えば列方向の情報をビット列として扱えますが、
欲しい情報によっては転置そのものを行わず、より高速に処理できる場合があります。

2.1 find_first

問題

$N\times N$ の $0$-$1$ 行列 $A$、長さ $N$ の整数列 $(p_0,p_1,\dots ,p_{N-1})$ が与えられる。ただし、$0\le p_j\le N-1$ とする。
各行は長さ $N$ のビット列として保持されているものとする。
各 $j$ $(0\le j\le N-1)$ に対して $A_{i,j}=1$ なる最小の $i(\ge p_j)$ を求めよ。存在しなければ、それを報告せよ。

解法

$i=0,1,\dots ,N-1$ の順に、答えが $i$ になるような $j$ $(0\le j\le N-1)$ を発見して、それを答えに反映することにします。

$i$ を処理する時点で、まだ答えが決まっておらず、かつ $p_j\le i$ であるような $j$ に対して $S_j=1$ となる $0$-$1$ 列 $S$ を持つことにします。最初は $S=0$ として、$i$ を処理する直前に、$p_j=i$ である全ての $j$ について $S_j=1$ とします。
$T=S\wedge A_i$ としましょう。$T_j=1$ であることは、$j$ の答えがまだ決まっておらず、$p_j\le i$ であり、かつ $A_{i,j}=1$ であることと同値であるため、$j$ に対する答えは $i$ であることを示しています。
従って、$T_j=1$ である全ての $j$ について答えを $i$ として、$S\leftarrow S\wedge \neg T$ とすればよいです。

$T$ のビットが立っている位置を $1$ つ探すのは $O\left(\frac{N}{w}\right)$ 時間でできます。（例えば、GCC なら std::bitset の _Find_first を使えます）各 $j$ について $T$ の $j$ ビット目が立つのは高々 $1$ 回ですから、$T$ のビットが立っている位置を探すパートは $O\left(\frac{N^2}{w}\right)$ 時間で可能です。$T$ を計算するのも、$N$ 回 $0$-$1$ 列のビット演算をするだけですから、結局 $O\left(\frac{N^2}{w}\right)$ 時間です。

2.2 popcount

問題

$N\times N$ の $0$-$1$ 行列 $A$ が与えられる。
各行は長さ $N$ のビット列として保持されているものとする。
各 $j=0,1,\dots ,N-1$ について、$A_{i,j}=1$ なる $i$ の個数を求めよ。

解法

各列の popcount だけが欲しいのであれば転置する必要はなく、Carry-Save Adder という手法で、より高速に解くことができます。

まずは、$A$ が $N\times w$ の $0$-$1$ 行列である場合を解きます。
$w$ 個の列の答えをそれぞれ独立に持つ代わりに、$w$ 列をまとめて管理してみましょう。

$l={\log }_{2}N$ として、$j=0,1,\dots ,l$ に対して、$w$ ビット整数を高々 $2$ 個持つ多重集合 ${\mathrm{bucket}}_0,{\mathrm{bucket}}_1,\dots ,{\mathrm{bucket}}_l$ を用意します。

$i=0,1,\dots ,N-1$ の順に $A_i$ を追加して、順次 $\mathrm{bucket}$ を更新することにします。$A_0,A_1,\dots ,A_{i-1}$ を追加した時点で、次の条件が満たされるようにしておきます。

• 各列 $k$ について、

$$\sum _{r=0}^{i-1}{A}_{r,k}=\sum _j{2}^j\left(\sum _{x\in {\mathrm{bucket}}_j}{x}_k\right)$$

が成立する。

• 各 $j$ について、$|{\mathrm{bucket}}_j|\le 2$ が成立する。

一つも追加していない場合は、${\mathrm{bucket}}_j=\varnothing$ とすれば条件を満たします。

$X=A_i$ を追加するとき、まず ${\mathrm{bucket}}_0\leftarrow {\mathrm{bucket}}_0\uplus \{X\}$ としたうえで、$\mathrm{modify}(0)$ を行います。
一般に $\mathrm{modify}(j)$ を次のように定義します。

$\mathrm{modify}(j)$

• $|{\mathrm{bucket}}_j|\le 2$ のとき

  • 何もせず終了します。

• $|{\mathrm{bucket}}_j|=3$ のとき

  • ${\mathrm{bucket}}_j$ の要素を $a,b,c$ とします。
  • $\mathrm{low}=a\oplus b\oplus c$ および $\mathrm{high}=(a\wedge b)\vee (b\wedge c)\vee (c\wedge a)$ とします。
  • ${\mathrm{bucket}}_j\leftarrow \{\mathrm{low}\}$ とし、${\mathrm{bucket}}_{j+1}\leftarrow {\mathrm{bucket}}_{j+1}\uplus \{\mathrm{high}\}$ として、$\mathrm{modify}(j+1)$ を行います。

アルゴリズムの正当性

まずは、 $\mathrm{modify}(j)$ によって、${\sum }_j{2}^j\left({\sum }_{x\in {\mathrm{bucket}}_j}{x}_k\right)$ が変化しないことを確認します。

一般に $x,y,z\in \{0,1\}$ に対して、

$$x+y+z=(x\oplus y\oplus z)+2((x\wedge y)\vee (y\wedge z)\vee (z\wedge x))$$

が成立します。つまり、各 $k$ $(0\le k<w)$ について、

$$a_k+b_k+c_k={\mathrm{low}}_k+2{\mathrm{high}}_k$$

が成立します。

よって、（$|{\mathrm{bucket}}_j|=3$ であった時に） ${\mathrm{bucket}}_j$ にあった $a,b,c$ による $2^j(a_k+b_k+c_k)$ という寄与は、$2^j{\mathrm{low}}_k+2^{j+1}{\mathrm{high}}_k$ に置き換わりますが、両者は等しいです。
したがって、$\mathrm{modify}(j)$ は ${\sum }_j{2}^j\left({\sum }_{x\in {\mathrm{bucket}}_j}{x}_k\right)$ を変えません。

$A_i$ を ${\mathrm{bucket}}_0$ に追加すると ${\sum }_j{2}^j\left({\sum }_{x\in {\mathrm{bucket}}_j}{x}_k\right)$ に $A_{i,k}$ が加わり、また $\mathrm{modify}(0)$ を行っても上に示した事実から${\sum }_j{2}^j\left({\sum }_{x\in {\mathrm{bucket}}_j}{x}_k\right)$ は変化しないため、処理後には

$$\sum _{r=0}^i{A}_{r,k}=\sum _j{2}^j\left(\sum _{x\in {\mathrm{bucket}}_j}{x}_k\right)$$

が成立しています。

また、確かに $\mathrm{modify}$ が終了した時点では全ての $j$ について $|{\mathrm{bucket}}_j|\le 2$ が成立しています。

最後に、$\mathrm{bucket}$ は ${\mathrm{bucket}}_l$ まで持っておけば十分であることを示します。
$S={\sum }_j{2}^j|{\mathrm{bucket}}_j|$ とします。$A_i$ まで追加した時点では $S=i+1$ です。また、$\mathrm{modify}$ の途中でも $S$ は変化しません。したがって処理中の任意の時点で $S\le N$ です。
よって、${\mathrm{bucket}}_j\ne \varnothing$ なら $2^j\le S\le N$ であり $j\le l$ が成立するため、${\mathrm{bucket}}_l$ まで持っておけば十分です。

以上から、このアルゴリズムが正しく動くことが分かりました。

計算量解析

計算量を解析します。$A_i$ を ${\mathrm{bucket}}_0$ に追加するたびに ${\sum }_j|{\mathrm{bucket}}_j|$ は $1$ 増加します。一方、$|{\mathrm{bucket}}_j|=3$ のときの $\mathrm{modify}(j)$ では、${\sum }_j|{\mathrm{bucket}}_j|$ は $1$ 減少します。

${\sum }_j|{\mathrm{bucket}}_j|$ は最初は $0$ であり、常に $0$ 以上です。また、増加する回数はちょうど $N$ 回なので、$|{\mathrm{bucket}}_j|=3$ の場合の処理が行われる回数は高々 $N$ 回です。

$\mathrm{modify}(j)$ が呼ばれるのは $\mathrm{modify}(j-1)$ によるものか、$A_i$ を追加した時に呼ばれる $\mathrm{modify}(0)$ のどちらかです。前者は先の議論により高々 $N$ 回、後者は明らかに $N$ 回ですから、$\mathrm{modify}$ が呼ばれる回数は高々 $2N$ 回です。

各更新で行われる演算はいずれも $w$ ビット整数のビット演算であるため、$O(1)$ 時間で実行できます。従って、上の処理全体は $O(N)$ 時間で行えます。

処理後は、各列 $k$ の答えを

$$\sum _{j=0}^l{2}^j\left(\sum _{x\in {\mathrm{bucket}}_j}{x}_k\right)$$

として復元できます。これは $O(wl)$ 時間で行えます。

以上から、$|A_i|=w$ の場合の計算量は $O(N+wl)=O(N+w\log N)$ です。

$|A_i|=N$ の場合は $N/w$ 個のブロックに分割してそれぞれにこのアルゴリズムを適用すれば、全体で $O\left(\frac{N^2}{w}+Nl\right)$ 時間のアルゴリズムが得られます。

$w=\Theta (\log N)$ という仮定の下では、$O\left(\frac{N^2}{w}\right)$ 時間です。

3. 適用例

https://atcoder.jp/contests/abc471/tasks/abc471_g

概要

$0,1,\dots ,K-1$ の各記号について母音かどうかが定められており、長さ $N$ の記号列 $A$ が与えられます。

各 $k=0,1,\dots ,K-1$ について、全ての $A_i$ を $(A_i+k)\mathrm{mod}K$ に置き換えたときの、母音からなる極大な連続区間の個数を求めよ。

• $1\le N\le 7\times 10^6$
• $1\le K\le 2300$

解法

母音を $1$、子音を $0$ に対応させます。
$0$-$1$ 列 $B$ の $1$ の極大区間の個数は $\frac{[B_0=1]+[B_{N-1}=1]+\#\{B_i\ne B_{i+1}\}}{2}$ です。第 $1$ 項、第 $2$ 項は容易に数えられるので、第 $3$ 項を考えます。

$0,1,\dots ,K-1$ が母音かどうかの $0$-$1$ 列を $V$ とし、$b_{i,k}=V_{(A_i+k)\mathrm{mod}K}$ として長さ $K$ の $0$-$1$ 列 $b_i$ を定義します。$b_i$ は $V$ の巡回シフトです。従って、$b_i\oplus b_{i+1}$ を $O\left(\frac{K}{w}\right)$ 時間で計算できます。

固定された $k$ に対して、第 $3$ 項は、ビット列 $b_0\oplus b_1,b_1\oplus b_2,\dots$ であって、$k$ ビット目が立っているものの個数と一致します。

従って、転置を用いるアルゴリズムにより $O\left(\frac{NK\log w}{w}\right)$ 時間、Carry-Save Adder を用いるアルゴリズムにより $O\left(\frac{NK}{w}+Kl\right)$ 時間で計算することができます。$w=\Theta (\log N)$ という仮定の下では、$O\left(\frac{NK}{w}\right)$ 時間です。

終わりに

明日の投稿者は@Alt--erさんです