Dynamic Wavelet Matrix

data_structure/dynamic_wavelet_matrix.hpp · GitHub で見る

共通ページ · 軽量実装なし

#include "lib/data_structure/dynamic_wavelet_matrix.hpp"

DynamicWaveletMatrix<int> wm(std::vector<int>{3, 1, 4, 1, 5});
wm.insert(2, -7);
wm.erase(0);
wm.set(1, 6);
auto value = wm.kth_smallest(0, wm.size(), 2);
auto frequency = wm.count(0, wm.size(), 1);

区間は \([l,r)\)、添字と kth_*・select の順位は 0 始まりです。kth_* の k は区間の長さ未満とします。

列長は INT_MAX 以下、挿入は 0 <= i <= size()、削除・代入・取得は 0 <= i < size()、区間は 0 <= l <= r <= size() とします。kth_* は 0 <= k < r-l が必要です。末尾の削除は非空の列に対して呼んでください。重複は個数分だけ数え、select(x,k) は元の列を左から数えた k 番目の出現位置を返します。空区間の個数・頻度は 0 です。

操作	内容・返り値	計算量
DynamicWaveletMatrix() / DynamicWaveletMatrix(a)	空／配列から構築	空: \(O(B)\)、配列: \(O(Bn)\)
build(a)	配列で再構築	\(O(Bn)\)
insert(i,x) / erase(i) / set(i,x)	挿入／削除／代入。erase は削除した値を返す	挿入・代入: 償却 \(O(B\log n)\)、削除: \(O(B\log n)\)
push_back(x) / pop_back()	末尾への挿入／削除。pop_back は削除した値を返す	追加: 償却 \(O(B\log n)\)、削除: \(O(B\log n)\)
access(i) / operator[](i)	要素の取得	\(O(B\log n)\)
count(l,r,x) / rank(x,r)	区間／\([0,r)\) にある x の個数。値域外なら 0	\(O(B\log n)\)
select(x,k)	0 始まり k 番目の x の位置。なければ -1	\(O(B\log n)\)
kth_smallest(l,r,k) / kth_largest(l,r,k)	小さい順／大きい順の 0 始まり k 番目	\(O(B\log n)\)
range_freq(l,r,upper)	upper 未満の個数	\(O(B\log n)\)
range_freq(l,r,lower,upper)	lower 以上 upper 未満の個数	\(O(B\log n)\)
prev_value(l,r,upper) / next_value(l,r,lower)	upper 未満の最大値／lower 以上の最小値。なければ nullopt	\(O(B\log n)\)
size() / empty()	要素数／空判定	\(O(1)\)
reserve(n) / clear()	領域予約／容量を残して空にする	reserve: \(O(BN)\)、clear: \(O(B)\)
memory_usage()	確保済みノード領域のバイト数	\(O(B)\)
swap(other)	内容を交換	\(O(B)\)

型 T は bool を除く 64 bit 以下の整数型です。DynamicWaveletMatrix<T,B> で値の bit 数を指定でき、既定は T の bit 数です。\(1\le B\le\mathrm{sizeof}(T)\cdot8\) とし、符号付き型の値域は \([-2^{B-1},2^{B-1})\)、符号なし型は \([0,2^B)\) です。格納する値はこの範囲に収めてください。

build(a) で再構築できます。reserve(n) は要素数の予約、memory_usage() は確保済みノード領域のバイト数です。clear() は容量を保持して空にします。copy・move・swap に対応し、move 元は空の列として再利用できます。

特定の値の個数には count(l,r,x) を使います。整数型の最大値も渡せます。range_freq(l,r,lower,upper) は lower >= upper なら 0 を返します。

prev_value(l,r,upper) は境界を含まず、next_value(l,r,lower) は境界を含みます。空区間や条件を満たす値がない場合は nullopt です。select(x,k) は x が値域外の場合や k が負の場合も -1 を返します。

include 前に DYNAMIC_QUAD_DISABLE_SIMD を定義すると SIMD・POPCNT・BMI2 経路を、DYNAMIC_QUAD_DISABLE_AVX512 を定義すると AVX-512 経路を無効にできます。

第3テンプレート引数に独自の BitVector を指定できます。move_rank と access_rank は省略できます。CPU 拡張命令をコンパイル時に指定する場合は、実行する CPU に合ったオプションを使ってください。

要素数を \(n\)、bit 数を \(B\) として、構築は \(O(Bn)\)、挿入・代入は償却 \(O(B\log n)\)、削除・クエリは \(O(B\log n)\) です。サイズ取得は \(O(1)\)、clear・memory_usage・swap は \(O(B)\)、reserve(N) は \(O(BN)\)、空間は \(O(BN)\) です。\(N\) は過去の最大要素数と予約容量の最大値です。

依存ライブラリ

• Dynamic Bit Vector
• Dynamic Quad Vector
• Dynamic Wavelet Matrix の4値の内部演算

検証コード

• test/dynamic_wavelet_matrix_kth.test.cpp

コード

#ifndef DYNAMIC_WAVELET_MATRIX_HPP
#define DYNAMIC_WAVELET_MATRIX_HPP

#include "../components/dynamic_bit_vector.hpp"
#include "../components/dynamic_quad_wavelet.hpp"
#include <algorithm>
#include <array>
#include <bit>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
#if defined(__AVX2__) && defined(__x86_64__)
#include <immintrin.h>
#endif

template <class T = std::int64_t, int Bits = sizeof(T) * 8, class BitVector = DynamicBitVector<>>
class DynamicWaveletMatrix {
	static_assert(std::is_integral_v<T> && !std::is_same_v<T, bool> && sizeof(T) <= 8);
	using U = std::make_unsigned_t<T>;
	static constexpr int width = std::numeric_limits<U>::digits;
	static_assert(1 <= Bits && Bits <= width);
	static constexpr U mask = std::numeric_limits<U>::max() >> (width - Bits);
	static constexpr U sign = std::is_signed_v<T> ? U(1) << (Bits - 1) : 0;
	static constexpr T minimum = std::is_signed_v<T> ? T(U(0) - sign) : T(0);
	static constexpr T maximum = T(mask >> std::is_signed_v<T>);
	std::array<BitVector, Bits> bits;
	std::array<int, Bits> middle{};
	int n = 0;
	static std::uint64_t pack_word(const U *values, int count, int shift) {
		std::uint64_t result = 0;
		int j = 0;
#if defined(__AVX2__) && defined(__x86_64__)
		if constexpr (sizeof(U) == 4 || sizeof(U) == 8) {
			__m128i amount = _mm_cvtsi32_si128(width - 1 - shift);
			if constexpr (sizeof(U) == 4) {
				for (; j + 8 <= count; j += 8) {
					__m256i value = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(values + j));
					result |= std::uint64_t(_mm256_movemask_ps(_mm256_castsi256_ps(_mm256_sll_epi32(value, amount)))) << j;
				}
			} else {
				for (; j + 4 <= count; j += 4) {
					__m256i value = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(values + j));
					result |= std::uint64_t(_mm256_movemask_pd(_mm256_castsi256_pd(_mm256_sll_epi64(value, amount)))) << j;
				}
			}
		}
#endif
		for (; j < count; ++j) result |= std::uint64_t((values[j] >> shift) & 1) << j;
		return result;
	}
	static bool contains(T value) { return minimum <= value && value <= maximum; }
	static U encode(T value) { return (U(value) & mask) ^ sign; }
	static T decode(U value) {
		value ^= sign;
		if constexpr (std::is_signed_v<T> && Bits < width)
			if (value & sign) value |= U(~mask);
		return T(value);
	}
	void check([[maybe_unused]] int l, [[maybe_unused]] int r) const { assert(0 <= l && l <= r && r <= n); }
	std::pair<int, int> interval(int l, int r, U key) const {
		for (int d = 0; d < Bits && l < r; ++d) {
			auto [a, b] = bits[d].rank1(l, r);
			if ((key >> (Bits - 1 - d)) & 1) l = middle[d] + a, r = middle[d] + b;
			else l -= a, r -= b;
		}
		return {l, r};
	}
	bool matches_suffix(int p, U key, int d) const {
		for (; d < Bits; ++d) {
			bool bit = (key >> (Bits - 1 - d)) & 1;
			auto [rank, value] = bits[d].access_rank(p);
			if (value != bit) return false;
			p = bit ? middle[d] + rank : p - rank;
		}
		return true;
	}

	static std::tuple<int, bool, int> move_rank(BitVector &row, int from, int to, bool value) {
		if constexpr (requires { row.move_rank(from, to, value); }) return row.move_rank(from, to, value);
		else {
			auto [rank, old] = row.erase_rank(from);
			int next = row.insert(to, value);
			return {rank, old, next};
		}
	}

  public:
	using value_type = T;
	DynamicWaveletMatrix() = default;
	explicit DynamicWaveletMatrix(const std::vector<T> &values) { build(values); }
	DynamicWaveletMatrix(const DynamicWaveletMatrix &) = default;
	DynamicWaveletMatrix &operator=(const DynamicWaveletMatrix &) = default;
	DynamicWaveletMatrix(DynamicWaveletMatrix &&other) noexcept { swap(other); }
	DynamicWaveletMatrix &operator=(DynamicWaveletMatrix &&other) noexcept {
		if (this != &other) clear(), swap(other);
		return *this;
	}
	void swap(DynamicWaveletMatrix &other) noexcept {
		bits.swap(other.bits);
		middle.swap(other.middle);
		std::swap(n, other.n);
	}
	int size() const { return n; }
	bool empty() const { return !n; }
	std::size_t memory_usage() const {
		std::size_t result = 0;
		for (const auto &row : bits) result += row.memory_usage();
		return result;
	}
	void reserve(int size) {
		for (auto &row : bits) row.reserve(size);
	}
	void clear() {
		for (auto &row : bits) row.clear();
		middle.fill(0);
		n = 0;
	}
	void build(const std::vector<T> &values) {
		assert(values.size() <= std::size_t(std::numeric_limits<int>::max()));
		n = int(values.size());
		std::vector<U> current(n), next(n);
		for (int i = 0; i < n; ++i) {
			assert(contains(values[i]));
			current[i] = encode(values[i]);
		}
		std::vector<std::uint64_t> packed((values.size() + 63) / 64);
		for (int d = 0; d < Bits; ++d) {
			int shift = Bits - 1 - d, zero = n;
			for (int i = 0; i < n;) {
				int end = std::min(64, n - i);
				std::uint64_t word = pack_word(current.data() + i, end, shift);
				packed[i >> 6] = word;
				zero -= std::popcount(word);
				i += end;
			}
			middle[d] = zero;
			bits[d].build(n, packed);
			if (d + 1 == Bits || !zero || zero == n) continue;
			std::array<int, 2> offset{0, zero};
			for (U value : current) next[offset[(value >> shift) & 1]++] = value;
			current.swap(next);
		}
	}
	T access(int p) const {
		assert(0 <= p && p < n);
		U result = 0;
		for (int d = 0; d < Bits; ++d) {
			auto [rank, value] = bits[d].access_rank(p);
			result = U((result << 1) | value);
			p = value ? middle[d] + rank : p - rank;
		}
		return decode(result);
	}
	T operator[](int p) const { return access(p); }
	void insert(int p, T value) {
		assert(0 <= p && p <= n && n < std::numeric_limits<int>::max() && contains(value));
		U key = encode(value);
		for (int d = 0; d < Bits; ++d) {
			bool bit = (key >> (Bits - 1 - d)) & 1;
			int rank = bits[d].insert(p, bit);
			p = bit ? middle[d] + rank : p - rank;
			middle[d] += !bit;
		}
		++n;
	}
	void push_back(T value) { insert(n, value); }
	T erase(int p) {
		assert(0 <= p && p < n);
		U result = 0;
		for (int d = 0; d < Bits; ++d) {
			auto [rank, value] = bits[d].erase_rank(p);
			result = U((result << 1) | value);
			p = value ? middle[d] + rank : p - rank;
			middle[d] -= !value;
		}
		--n;
		return decode(result);
	}
	T pop_back() { return erase(n - 1); }
	void set(int p, T value) {
		assert(0 <= p && p < n && contains(value));
		int old = p;
		U key = encode(value);
		for (int d = 0; d < Bits; ++d) {
			bool bit = (key >> (Bits - 1 - d)) & 1;
			int z = middle[d];
			if ((!z && bit) || (z == n && !bit)) continue;
			if (old == p) {
				auto [rank, previous] = bits[d].set_rank(old, bit);
				old = previous ? z + rank : old - rank;
				middle[d] += int(previous) - int(bit);
				p = bit ? middle[d] + rank : p - rank;
			} else {
				auto [rank, previous, next] = move_rank(bits[d], old, p, bit);
				old = previous ? z + rank : old - rank;
				middle[d] -= !previous;
				p = bit ? middle[d] + next : p - next;
				middle[d] += !bit;
			}
		}
	}
	int count(int l, int r, T value) const {
		check(l, r);
		if (l == r || !contains(value)) return 0;
		U key = encode(value);
		int d = 0;
		for (; d < Bits; ++d) {
			int zero = middle[d];
			if (zero && zero != n) break;
			if (bool((key >> (Bits - 1 - d)) & 1) != !zero) return 0;
		}
		if (r - l == 1) return matches_suffix(l, key, d);
		for (; d < Bits && l < r; ++d) {
			auto [a, b] = bits[d].rank1(l, r);
			if ((key >> (Bits - 1 - d)) & 1) l = middle[d] + a, r = middle[d] + b;
			else l -= a, r -= b;
		}
		return r - l;
	}
	int rank(T value, int r) const { return count(0, r, value); }
	int range_freq(int l, int r, T upper) const {
		check(l, r);
		if (upper <= minimum) return 0;
		if (maximum < upper) return r - l;
		U key = encode(upper);
		int result = 0;
		for (int d = 0; d < Bits && l < r; ++d) {
			auto [a, b] = bits[d].rank1(l, r);
			if ((key >> (Bits - 1 - d)) & 1) {
				result += r - l - b + a;
				l = middle[d] + a, r = middle[d] + b;
			} else l -= a, r -= b;
		}
		return result;
	}
	int range_freq(int l, int r, T lower, T upper) const {
		check(l, r);
		if (l == r || !(lower < upper)) return 0;
		if (lower <= minimum) return range_freq(l, r, upper);
		if (maximum < upper) return r - l - range_freq(l, r, lower);
		U low = encode(lower), high = encode(upper);
		int common = Bits - std::bit_width(U(low ^ high)), d = 0;
		for (; d < common && l < r; ++d) {
			auto [a, b] = bits[d].rank1(l, r);
			if ((low >> (Bits - 1 - d)) & 1) l = middle[d] + a, r = middle[d] + b;
			else l -= a, r -= b;
		}
		if (l == r) return 0;
		auto [a, b] = bits[d].rank1(l, r);
		int left = l - a, right = r - b, result = right - left;
		for (int level = d + 1; level < Bits && left < right; ++level) {
			auto [x, y] = bits[level].rank1(left, right);
			if ((low >> (Bits - 1 - level)) & 1) {
				result -= right - left - y + x;
				left = middle[level] + x, right = middle[level] + y;
			} else left -= x, right -= y;
		}
		left = middle[d] + a, right = middle[d] + b;
		for (int level = d + 1; level < Bits && left < right; ++level) {
			auto [x, y] = bits[level].rank1(left, right);
			if ((high >> (Bits - 1 - level)) & 1) {
				result += right - left - y + x;
				left = middle[level] + x, right = middle[level] + y;
			} else left -= x, right -= y;
		}
		return result;
	}
	T kth_smallest(int l, int r, int k) const {
		check(l, r);
		assert(0 <= k && k < r - l);
		U result = 0;
		for (int d = 0; d < Bits; ++d) {
			if (r - l == 1) {
				for (; d < Bits; ++d) {
					auto [rank, bit] = bits[d].access_rank(l);
					result = U((result << 1) | bit);
					l = bit ? middle[d] + rank : l - rank;
				}
				break;
			}
			auto [a, b] = bits[d].rank1(l, r);
			int zero = r - l - b + a;
			bool value = k >= zero;
			result = U((result << 1) | value);
			if (value) k -= zero, l = middle[d] + a, r = middle[d] + b;
			else l -= a, r -= b;
		}
		return decode(result);
	}
	T kth_largest(int l, int r, int k) const {
		check(l, r);
		assert(0 <= k && k < r - l);
		return kth_smallest(l, r, r - l - 1 - k);
	}
	std::optional<T> prev_value(int l, int r, T upper) const {
		int k = range_freq(l, r, upper);
		return k ? std::optional<T>(kth_smallest(l, r, k - 1)) : std::nullopt;
	}
	std::optional<T> next_value(int l, int r, T lower) const {
		int k = range_freq(l, r, lower);
		return k < r - l ? std::optional<T>(kth_smallest(l, r, k)) : std::nullopt;
	}
	int select(T value, int k) const {
		if (k < 0 || !contains(value)) return -1;
		U key = encode(value);
		int l = 0, r = n, d = 0;
		for (; d < Bits && l < r; ++d) {
			if constexpr (requires { bits[d].access_rank(l); }) {
				if (r - l == 1) {
					if (k || !matches_suffix(l, key, d)) return -1;
					break;
				}
			}
			auto [a, b] = bits[d].rank1(l, r);
			if ((key >> (Bits - 1 - d)) & 1) l = middle[d] + a, r = middle[d] + b;
			else l -= a, r -= b;
		}
		if (k >= r - l) return -1;
		int p = l + k;
		for (--d; d >= 0; --d) {
			bool bit = (key >> (Bits - 1 - d)) & 1;
			p = bits[d].select(bit, p - (bit ? middle[d] : 0));
		}
		return p;
	}
};

template<class T, int Bits>
class DynamicWaveletMatrix<T, Bits, DynamicBitVector<>> : public dynamic_wavelet_detail::QuadMatrix<T, Bits> {
	using Base = dynamic_wavelet_detail::QuadMatrix<T, Bits>;
  public:
	using Base::Base;
	void swap(DynamicWaveletMatrix& other) noexcept { Base::swap(other); }
};

#endif

Back to top page