Vertical Code + Succinct Bit Vector + Compressed Suffix Array

前回までの記事では、完全一致検索を高速化するためにリファレンス文字列に加えてSuffix Arrayを保持していました。しかし、文字列自体よりもSuffix Arrayのほうが大きいのが問題でした。この記事ではSuffix Arrayを圧縮して保持する方法の一つを説明します。文字列そのものよりも小さな追加データ(Compressed Suffix Array, CSA)を用いて、引数iを受け取ってsuffix_array[i]を返す機能を作れます。そのための基礎としてVertical CodeとSuccinct Bit Vectorが必要になりますが、この記事ではそれらについても紹介します。

d.hatena.ne.jp

岡野原大輔, 現実的な圧縮付全文索引

これらの記事を参考にしました。後者の論文の3.1節の最初において、CSAの方法としてGrossiらの論文とSadakaneらの論文がreferされています。後者の論文自身はSadakaneらの論文に基づいています。一方前者のブログ記事はGrossiらの論文に基づいています。今回の記事ではSadakaneらの論文の方法に基づいています。

Vertical Code

現実的な圧縮付全文索引、Vertical Codeを調べたよ - EchizenBlog-Zweiで説明されています。

以下のコード例を示します。popcount関数は引数の立っているビットを数えるもので、SSE4.2以降にpopcntという等価な命令がありますが今回は基本演算で実装しました。トリッキーな実装になっていますが最速ではありません。cf. 2009-07-04 - 当面C#と.NETな記録

ブロックサイズが64で決め打ちされていますが、これは簡単のためであり、ブロックサイズ64で性能が最善になるとは限りません。

class VerticalCode {
private:

	static constexpr int block = 64;
	uint64_t size;
	std::vector<uint64_t>psi;
	std::vector<uint64_t>code;
	std::vector<uint64_t>code_index;

	static uint64_t popcount(uint64_t x) {
		x = (x & 0x5555555555555555ULL) + ((x & 0xAAAAAAAAAAAAAAAAULL) >> 1);
		x = (x & 0x3333333333333333ULL) + ((x & 0xCCCCCCCCCCCCCCCCULL) >> 2);
		x = (x & 0x0F0F0F0F0F0F0F0FULL) + ((x & 0xF0F0F0F0F0F0F0F0ULL) >> 4);
		x *= 0x0101010101010101ULL;
		return x >> 56;
	}

public:

	VerticalCode(const std::vector<uint64_t>& input) {
		//inputが元の数列だとして、差分化してvertical codeで保持する。
		//inputの各要素は[0,N)の並べ替えであることを仮定する。

		{
			auto testarray = input;
			sort(testarray.begin(), testarray.end());
			for (int i = 0; i < testarray.size(); ++i)assert(testarray[i] == i);
		}

		size = input.size();
		const int64_t code_length = (input.size() + block - 1) / block;

		//psi[i]には、0～(i-1)ブロック内の差分列の総和をsizeで割った余りを格納する。
		//「差分列の総和」は元の数列の値とは異なる。なぜなら、
		//ここで言う「差分列」は、圧縮しやすい単調増加列にするため、差分に-1や+Nしているからである。
		psi = std::vector<uint64_t>(code_length, 0);
		//for (int64_t i = 1; i < code_length; ++i)psi[i] = input[i * block - 1];

		code.clear();
		code_index = std::vector<uint64_t>(code_length + 1, 0);

		for (int64_t b = 0; b < code_length; ++b) {
			const int64_t offset = b * block;

			//offset番目からblock個について、単調増加な差分列を作る。
			std::vector<uint64_t>diff(std::min(int64_t(64), int64_t(input.size() - offset)), 0);
			uint64_t diffmax = 0;
			for (int64_t i = 0; i < diff.size(); ++i) {
				const int64_t j = offset + i;
				if (j == 0)diff[i] = input[j];
				else if (input[j - 1] < input[j])diff[i] = input[j] - input[j - 1] - 1;
				else {
					//↓はinputが[0,N)の並べ替えであれば必ず成り立つ。
					//assert(input[j - 1] - input[j] < input.size());
					diff[i] = input[j] - input[j - 1] - 1 + input.size();
				}
				diffmax = std::max(diffmax, diff[i]);
			}
			const int64_t bit_demand = SparseTableMin<uint64_t>::log2_ceiling(diffmax + 1);
			
			code.resize(code.size() + bit_demand);
			code_index[b + 1] = code.size();
			for (int i = 0; i < diff.size(); ++i) {
				for (uint64_t num = diff[i], j = code_index[b]; num; num /= 2, ++j) {
					code[j] |= (num & uint64_t(1)) << i;
				}
			}

			if (b < code_length - 1) {
				psi[b + 1] = psi[b];
				for (const int x : diff)psi[b + 1] += x;
				psi[b + 1] %= size;
			}
		}
	}

	VerticalCode() {
		VerticalCode(std::vector<uint64_t>{0});
	}

	uint64_t Get(const int64_t index) {
		const int64_t b = index / block;
		const auto bits = index - (b * block);

		const uint64_t bitmask = 0xFFFF'FFFF'FFFF'FFFFULL >> (63 - bits);

		uint64_t answer = psi[b];
		for (uint64_t i = code_index[b]; i < code_index[b + 1]; ++i) {
			answer += popcount(code[i] & bitmask) << (i - code_index[b]);
		}
		return (answer + index) % size;
	}
};

Succinct Bit Vector

長さNのビット列に対して、少ない追加データを用意することで、rank操作とselect操作という2種類の操作を定数時間で処理できるようにしたものをSuccinct Bit Vectorと呼びます。曖昧な言い方になっています。実は、長さNのビット列に対してo(N)の追加データによりrankとselectを定数時間で処理できるような方法が存在しますが、以下のコード例はそうなっていません。具体的には、以下のコードでは256bitごとと32bitごとのテーブルが用意されますが、o(N)を実現するためにはこの値を(256と32の決め打ちではなく)Nに応じた適切な値とする必要があります。Bit Vectorを実際に使いたい人は有名なライブラリがあるのでそれを使いましょう。あとCompressed Suffix Arrayではselect操作は構築のときにのみ使われて、クエリ処理のときには使われません。そのため以下のコード例ではselectが単純で低速な実装になっています。

class BruteForceBitVector {
private:

	std::vector<uint64_t>B;

public:

	BruteForceBitVector(const std::vector<uint64_t>& input) {
		B = input;
	}
	BruteForceBitVector() {
		BruteForceBitVector(std::vector<uint64_t>{0});
	}

	bool access(const uint64_t index) {
		return (B[index / 64] & (uint64_t(1) << (index % 64))) != 0;
	}

	uint64_t rank(const uint64_t index) {
		//[0,index)にいくつ立っているビットがあるか求めて返す。
		uint64_t answer = 0;
		for (int i = 0; i < index; ++i)if (access(i))++answer;
		return answer;
	}

	uint64_t select1(const uint64_t count) {
		//count番目(0-origin)の1のビットの位置を求めて返す。低速。
		uint64_t acc = 0;
		for (int i = 0; i < B.size() * 64; ++i)if (access(i))if (acc++ == count)return i;
		return std::numeric_limits<uint64_t>::max();
	}
};
class BitVector {
private:

	std::vector<uint64_t>B;
	std::vector<uint64_t>rank_global_256;
	std::vector<uint8_t>rank_local_64;

public:

	BitVector(const std::vector<uint64_t>& input) {

		B = input;

		//rank_global_256[i]の末尾以外には、[0,256*i)の範囲に立っているビットの数を格納する。
		//末尾にはデータ全体で立っているビットの数を格納する。
		rank_global_256.resize(((input.size() + 3) / 4) + 1);

		//rank_local_64[i]には、[256*(i/4),64*i)の範囲に立っているビットの数を格納する。
		rank_local_64.resize(input.size());

		uint64_t global_sum = 0;
		uint64_t local_sum = 0;
		for (int64_t i = 0; i < input.size(); ++i) {
			if (i % 4 == 0) {
				local_sum = 0;
			}
			rank_local_64[i] = local_sum;
			const uint64_t pop = popcount(input[i]);
			global_sum += pop;
			local_sum += pop;
			if (i % 4 == 3) {
				rank_global_256[i / 4 + 1] = global_sum;
			}
		}
		rank_global_256.back() = global_sum;
	}

	BitVector() {
		BitVector(std::vector<uint64_t>{0});
	}

	BitVector(const BitVector& obj) {
		this->B = obj.B;
		this->rank_global_256 = obj.rank_global_256;
		this->rank_local_64 = obj.rank_local_64;
	}

	bool access(const uint64_t index) {
		return (B[index / 64] & (uint64_t(1) << (index % 64))) != 0;
	}

	uint64_t rank(const uint64_t index) {
		//[0,index)にいくつ立っているビットがあるか求めて返す。
		uint64_t answer = rank_global_256[index / 256] + rank_local_64[index / 64];
		answer += popcount(B[index / 64] & ((uint64_t(1) << (index % 64)) - 1));
		return answer;
	}

	uint64_t select1(const uint64_t count) {
		//count番目(0-origin)の1のビットの位置を求めて返す。低速。

		int64_t L = 0, R = B.size() * 64;
		if (rank_global_256.back() <= count)return std::numeric_limits<uint64_t>::max();
		if (rank_global_256.back() == count - 1 && access(R - 1))return R - 1;

		//欲しい位置は[L,R)のどこかだとする。
		while (L + 1 < R) {
			const int64_t mid = (L + R) / 2;
			if (rank(mid) <= count)L = mid;
			else R = mid;
		}
		return L;
	}

	uint64_t sum1() { return rank_global_256.back(); }

	static uint64_t popcount(uint64_t x) {
		x = (x & 0x5555555555555555ULL) + ((x & 0xAAAAAAAAAAAAAAAAULL) >> 1);
		x = (x & 0x3333333333333333ULL) + ((x & 0xCCCCCCCCCCCCCCCCULL) >> 2);
		x = (x & 0x0F0F0F0F0F0F0F0FULL) + ((x & 0xF0F0F0F0F0F0F0F0ULL) >> 4);
		x *= 0x0101010101010101ULL;
		return x >> 56;
	}
};

Compressed Suffix Array

下記のコードでは、現実的な圧縮付全文索引で説明されているアルゴリズムが実装されています。SAISSuffixArrayConstructionは前回の記事のコードを想定しています。

template<uint64_t compress_order>class CompressedSuffixArray {

	//2^compress_order個ごとにサンプルする。

private:

	VerticalCode psi_vertical_code;
	std::vector<uint64_t>suffix_array_sumple;
	BitVector B;
	int64_t suffix_array_size;

public:

	CompressedSuffixArray(const std::vector<uint8_t>& input) {

		assert(1 <= compress_order && compress_order <= 10);
		for (int64_t i = 0; i < input.size() - 1; ++i)assert(input[i] != 0);
		assert(input.back() == 0);

		const auto suffix_array = SAISSuffixArrayConstruction(input);
		suffix_array_size = suffix_array.size();

		//Inverse Suffix Arrayを作る。
		std::vector<int>inverse_suffix_array(suffix_array.size(), 0);
		for (int64_t i = 0; i < inverse_suffix_array.size(); ++i) {
			inverse_suffix_array[suffix_array[i]] = i;
		}

		//Inverse Suffix Arrayを使ってPsi Arrayを作り、Vertical Codeにする。
		std::vector<uint64_t>psi(suffix_array.size(), 0);
		for (int64_t i = 0; i < psi.size(); ++i) {
			psi[i] =
				(suffix_array[i] == suffix_array.size() - 1) ?
				inverse_suffix_array[0] :
				inverse_suffix_array[suffix_array[i] + 1];
		}
		psi_vertical_code = VerticalCode(psi);

		//ビット列Bを作る。定義は
		//Bのi番目のビットがtrue⇔SA[i]が2^compress_orderで割り切れる
		//とする。(元の岡野原論文の定義)
		//加えてSA[i]==SA.size()-1のときもサンプルを保存しておく。
		//そうすることでGetのときに剰余を取る必要がなくなる。
		const auto interval = uint64_t(1) << compress_order;
		std::vector<uint64_t>b_tmp((suffix_array.size() + 63) / 64);
		for (int i = 0; i < suffix_array.size(); ++i) {
			if ((suffix_array[i] % interval) == 0 || suffix_array[i] == suffix_array.size() - 1) {
				b_tmp[i / 64] |= uint64_t(1) << (i % 64);
			}
		}
		B = BitVector(b_tmp);

		//Bのビットが立っている位置のsuffix arrayだけをサンプリングして保存しておく。
		const int64_t sample_size = B.sum1();
		suffix_array_sumple = std::vector<uint64_t>(sample_size, 0);
		for (int64_t i = 0; i < sample_size; ++i) {
			suffix_array_sumple[i] = suffix_array[B.select1(i)];
		}
	}

	uint64_t Get(int64_t index) {

		//suffix_array[index]を返す。

		for(int64_t move_count = 0;;++move_count) {

			if (B.access(index)) {

				//剰余について
				//この記述の中ではk=2^compress_order, t=move_countとする。
				//元の岡野原論文の定義だと、B.access(index)がtrueになる⇔SA[index]がkで割り切れる　であった。
				//B.access(index)がfalseの場合、index:=psi[index]とすることで、"返り値が1増える"効果があるので、
				//たかだかk-1回繰り返しpsiを作用させることで、B.access(index)がtrueになる。
				//そのとき、psiを作用させた回数をtとすると、SA_sample[B.rank(index)]-tを返せば良いと書かれていた。
				//しかし実際には、SA.size()がkで割り切れないとき、SA_sample[B.rank(index)]==0になることがある。
				//そのため、元論文の通りSA[index]がkで割り切れるもののみサンプリングしたなら、
				//元論文と違いSA.size()で割った余りを返す必要がある。
				//今回の実装では元論文と違いSA[index]==SA.size()-1のときもサンプリングしているので、
				//tが非ゼロのときSA_sample[B.rank(index)]==0にはならず、剰余の処理は必要ない。

				return (int64_t(suffix_array_sumple[B.rank(index)]) - move_count);
				//return (int64_t(suffix_array_sumple[B.rank(index)]) - move_count + suffix_array_size) % suffix_array_size;
			}
			index = psi_vertical_code.Get(index);
		}

		assert(0);
		return 0;
	}
};