Wavelet Matrix の座標圧縮

components/wavelet_compression.hpp · GitHub で見る

共通ページ

通常の Wavelet Matrix と重み付き Wavelet Matrix が共有する内部の座標圧縮です。Wavelet Matrix 本体やビット列には依存しません。

wavelet_detail::ValueIndex<T>::build(a) は元の値を昇順・重複なしにした values と、各位置の圧縮値 original を作ります。code(x) は values 内で x 未満の値の数を返し、x が登録されていなくても使えます。

構築は整数型向けの基数ソートと一般型向けの比較ソートを使い分けます。圧縮関数 compress は、通常版が利用する値ごとの出現位置と区切りも生成できます。compress<T,false> はそれらを生成せず、ValueIndex はこの経路を使います。出現位置を生成する呼び出しでは、前回の出力を初期化してから処理するため、同じ格納先を再利用できます。

要素数を \(n\)、異なる値の数を \(\sigma\) として、構築は整数型では \(O(n)\)、一般型では \(O(n\log n)\) 時間、空間は \(O(n)\) です。code は \(O(\log\sigma)\) 時間・\(O(1)\) 追加空間です。整数型のビット幅を定数とします。

検証コード

• test/offline_range_kth_smallest.test.cpp
• test/weighted_wavelet_matrix_rectangle.test.cpp

コード

#ifndef WAVELET_COMPRESSION_HPP
#define WAVELET_COMPRESSION_HPP

#include <algorithm>
#include <array>
#include <bit>
#include <cassert>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <vector>

namespace wavelet_detail {
template<class T>
std::uint32_t lower_code(const std::vector<T> &values, bool dense, const T &x) {
	if (values.empty()) return 0;
	if constexpr (std::is_integral_v<T> && sizeof(T) <= 8)
		if (dense) {
			if (x <= values.front()) return 0;
			if (values.back() < x) return static_cast<int>(values.size());
			return std::uint64_t(x) - std::uint64_t(values.front());
		}
	int p = 0, length = static_cast<int>(values.size());
	while (length > 1) {
		int half = length / 2;
		p += (values[p + half] < x) * half;
		length -= half;
	}
	return p + (values[p] < x);
}

template<class T, bool WithPositions = true>
std::vector<std::uint32_t> compress(const std::vector<T> &a, std::vector<T> &values,
	int &ordered, std::vector<int> &positions, std::vector<int> &offsets) {
	assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
	const int n = static_cast<int>(a.size());
	ordered = 0;
	if constexpr (WithPositions) offsets.clear();
	std::vector<std::uint32_t> current(n);
	bool ascending = std::is_sorted(a.begin(), a.end());
	if (ascending || std::is_sorted(a.rbegin(), a.rend())) {
		ordered = ascending ? 1 : -1;
		values.clear();
		values.reserve(n);
		for (int i = 0; i < n; ++i) {
			int j = ascending ? i : n - 1 - i;
			if (values.empty() || values.back() != a[j]) values.push_back(a[j]);
			current[j] = static_cast<int>(values.size()) - 1;
		}
		if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
		if constexpr (WithPositions) positions.clear();
		return current;
	}
	if constexpr (std::is_integral_v<T> && !std::is_same_v<T, bool> && sizeof(T) <= 8) {
		if (n >= 1024) {
			using U = std::make_unsigned_t<T>;
			struct [[gnu::packed]] Item {
				U key;
				std::uint32_t index;
			};
			std::vector<Item> order(n);
			U lo = std::numeric_limits<U>::max(), hi = 0;
			constexpr U sign = std::is_signed_v<T> ? U(1) << (sizeof(T) * 8 - 1) : U(0);
			for (int i = 0; i < n; ++i) {
				U key = U(a[i]) ^ sign;
				order[i] = {key, std::uint32_t(i)};
				lo = std::min(lo, key);
				hi = std::max(hi, key);
			}
			if (hi - lo < 65536 && std::uint64_t(hi - lo) < unsigned(n)) {
				std::vector<int> codes(std::size_t(hi - lo) + 1, -1);
				for (const auto &v : order) codes[v.key - lo] = 0;
				values.clear();
				values.reserve(codes.size());
				for (std::size_t i = 0; i < codes.size(); ++i)
					if (codes[i] == 0) {
						codes[i] = static_cast<int>(values.size());
						values.push_back(T(U(lo + U(i)) ^ sign));
					}
				for (const auto &v : order) current[v.index] = codes[v.key - lo];
				if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
				if constexpr (WithPositions) positions.clear();
				return current;
			}
			std::vector<Item> next(n);
			int width = std::bit_width(U(lo ^ hi)), span_width = std::bit_width(U(hi - lo));
			if ((width + 10) / 11 > (span_width + 10) / 11) {
				for (auto &v : order) v.key -= lo;
				width = span_width;
			}
			for (int shift = 0; shift < width; shift += 11) {
				std::array<int, 2048> offsets{};
				for (const auto &v : order) ++offsets[(v.key >> shift) & 2047];
				int total = 0;
				for (int &c : offsets) {
					int t = c;
					c = total;
					total += t;
				}
				for (const auto &v : order) next[offsets[(v.key >> shift) & 2047]++] = v;
				order.swap(next);
			}
			values.clear();
			values.reserve(n);
			U previous = 0;
			if constexpr (WithPositions) {
				positions.resize(n);
				offsets.reserve(std::size_t(n) + 1);
			}
			for (int j = 0; j < n; ++j) {
				const auto &v = order[j];
				if (values.empty() || v.key != previous) {
					values.push_back(a[v.index]);
					if constexpr (WithPositions) offsets.push_back(j);
				}
				previous = v.key;
				current[v.index] = static_cast<int>(values.size()) - 1;
				if constexpr (WithPositions) positions[j] = v.index;
			}
			if constexpr (WithPositions) offsets.push_back(n);
			if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
			return current;
		}
	}
	values = a;
	std::sort(values.begin(), values.end());
	values.erase(std::unique(values.begin(), values.end()), values.end());
	if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
	for (int i = 0; i < n; ++i) current[i] = lower_code<T>(values, false, a[i]);
	if constexpr (WithPositions) positions.clear();
	return current;
}

template<class T> struct ValueIndex {
	int n = 0;
	std::vector<T> values;
	std::vector<std::uint32_t> original;
	void build(const std::vector<T> &a) {
		assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
		n = static_cast<int>(a.size());
		int ordered = 0;
		std::vector<int> positions, offsets;
		original = compress<T, false>(a, values, ordered, positions, offsets);
	}
	std::uint32_t code(const T &x) const { return lower_code(values, false, x); }
};
}

#endif

Back to top page