cp-library

競技プログラミング用 C++ ライブラリ

View the Project on GitHub qandori/cp-library

:warning: Point Set Wavelet Matrix (クエリ先読み)

data_structure/point_set_wavelet_matrix.hpp · GitHub で見る

共通ページ

#include "lib/data_structure/point_set_wavelet_matrix.hpp"

std::vector<int> a{3, 1, 4, 1, 5};
std::vector<std::pair<int, int>> candidates{ {0, 10}, {2, 7} };
PointSetWaveletMatrix<int> wm(a, candidates);
wm.set(0, 10);
auto answer = wm.min_count_largest(0, 5, 12);

配列の一点代入と、区間内の順位・和のクエリを扱います。初期値以外に代入する全ての (位置,値) の事前登録が必要です。通常は構築前に更新クエリを先読みし、候補を candidates に集めます。初期値は自動登録されます。候補が最初から分かっていれば、更新クエリ自体の先読みは不要です。順位・和を求めるクエリは先読みせずに処理できます。区間は \([l,r)\)、順位は 0 始まりです。

一点操作と候補の位置は 0 <= i < size()、区間は 0 <= l <= r <= size() とします。kth_* は 0 <= k < r-l、sum_smallest・sum_largest は 0 <= k <= r-l が必要です。和の k は個数で、0 個の和は 0 です。重複する要素は個数分だけ数え、空区間の個数・和は 0 です。同じ候補を複数登録しても一つにまとめます。列長と異なる登録 (位置,値) の総数はそれぞれ INT_MAX 以下とします。入力は変更せず、構築後に破棄できます。

操作 内容・返り値 計算量
PointSetWaveletMatrix() 空で構築 \(O(1)\)
PointSetWaveletMatrix(a,candidates = {}) / build(a,candidates = {}) 初期配列と代入候補から構築/再構築 \(O(m\log m)\)
size() / empty() 配列の長さ/空判定 \(O(1)\)
access(i) / operator[](i) 現在の値 \(O(1)\)
contains(i,x) 位置 i に値 x を代入できるか \(O(\log m)\)
set(i,x) 一点代入。未登録なら false を返し、内容を保持 \(O(\log m\log\sigma)\)
count(l,r,x) / rank(x,r) 区間/\([0,r)\) の x の個数 \(O(\log m\log\sigma)\)
range_freq(l,r,upper) upper 未満の個数 \(O(\log m\log\sigma)\)
range_freq(l,r,lower,upper) lower 以上 upper 未満の個数 \(O(\log m\log\sigma)\)
range_sum(l,r) 区間和 \(O(\log m)\)
range_sum(l,r,upper) upper 未満の値の和 \(O(\log m\log\sigma)\)
range_sum(l,r,lower,upper) 指定した値域の和 \(O(\log m\log\sigma)\)
kth_smallest(l,r,k) / kth_largest(l,r,k) 小さい順/大きい順の k 番目 \(O(\log m\log\sigma)\)
sum_smallest(l,r,k) / sum_largest(l,r,k) 下位/上位 k 個の和 \(O(\log m\log\sigma)\)
min_count_smallest(l,r,target) / min_count_largest(l,r,target) 各順序で和が target 以上になる最小個数。届かなければ -1 \(O(\log m\log\sigma)\)
max_count_smallest(l,r,budget) / max_count_largest(l,r,budget) 各順序で和が budget 以下の最大個数。空の選択も不可能なら -1 \(O(\log m\log\sigma)\)

T は bool を除く 64 bit 以下の整数型、Sum は符号付き整数型です。PointSetWaveletMatrix<T,Sum> で和の型を指定できます。既定は 32 bit 以下の値なら std::int64_t、64 bit の値なら __int128_t です。中間の部分和も Sum に収めてください。

和の境界探索 min_count_*・max_count_* には初期値と登録した候補全てが非負であることが必要です。前者は target <= 0 なら 0、後者は budget < 0 なら -1 を返します。その他のクエリは負の値も扱えます。

異なる登録 (位置,値) の総数を \(m\)、異なる値の数を \(\sigma\) として、構築は \(O(m\log m)\) 時間、空間は \(O(m\log\sigma)\) です。更新と区間の順位・和の探索は \(O(\log m\log\sigma)\)、値域を指定しない区間和は \(O(\log m)\)、候補の確認は \(O(\log m)\)、サイズと要素取得は \(O(1)\) です。

依存ライブラリ

コード

#ifndef POINT_SET_WAVELET_MATRIX_HPP
#define POINT_SET_WAVELET_MATRIX_HPP

#include "weighted_wavelet_matrix.hpp"

template <class T = std::int64_t, class Sum = std::conditional_t<(sizeof(T) < 8), std::int64_t, __int128_t>>
class PointSetWaveletMatrix {
	std::vector<T> keys;
	std::vector<int> offsets{0}, current;
	WeightedWaveletMatrix<T, int, Sum> wm;
	void check([[maybe_unused]] int l, [[maybe_unused]] int r) const { assert(0 <= l && l <= r && r <= size()); }
	int find(int p, const T &x) const {
		assert(0 <= p && p < size());
		return std::lower_bound(keys.begin() + offsets[p], keys.begin() + offsets[p + 1], x) - keys.begin();
	}

  public:
	using value_type = T;
	using sum_type = Sum;
	PointSetWaveletMatrix() = default;
	explicit PointSetWaveletMatrix(const std::vector<T> &a, const std::vector<std::pair<int, T>> &candidates = {}) { build(a, candidates); }
	void build(const std::vector<T> &a, const std::vector<std::pair<int, T>> &candidates = {}) {
		assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
		std::vector<std::pair<int, T>> points = candidates;
		points.reserve(a.size() + candidates.size());
		for (int i = 0; i < static_cast<int>(a.size()); ++i) points.emplace_back(i, a[i]);
		std::sort(points.begin(), points.end());
		points.erase(std::unique(points.begin(), points.end()), points.end());
		assert(points.size() <= std::size_t(std::numeric_limits<int>::max()));
		offsets.assign(a.size() + 1, 0);
		keys.clear();
		keys.reserve(points.size());
		for (const auto &[p, x] : points) {
			assert(0 <= p && p < static_cast<int>(a.size()));
			++offsets[p + 1];
			keys.push_back(x);
		}
		for (int i = 0; i < static_cast<int>(a.size()); ++i) offsets[i + 1] += offsets[i];
		current.resize(a.size());
		std::vector<int> weights(points.size());
		for (int i = 0; i < size(); ++i) weights[current[i] = find(i, a[i])] = 1;
		wm.build(keys, weights);
	}
	int size() const { return static_cast<int>(current.size()); }
	bool empty() const { return current.empty(); }
	T access(int p) const { assert(0 <= p && p < size()); return keys[current[p]]; }
	T operator[](int p) const { return access(p); }
	bool contains(int p, const T &x) const {
		int q = find(p, x);
		return q < offsets[p + 1] && keys[q] == x;
	}
	bool set(int p, const T &x) {
		int q = find(p, x);
		if (q == offsets[p + 1] || keys[q] != x) return false;
		wm.move_weight(current[p], q, 1);
		current[p] = q;
		return true;
	}
	int count(int l, int r, const T &x) const { check(l, r); return wm.count(offsets[l], offsets[r], x); }
	int rank(const T &x, int r) const { return count(0, r, x); }
	int range_freq(int l, int r, const T &upper) const { check(l, r); return wm.range_freq(offsets[l], offsets[r], upper); }
	int range_freq(int l, int r, const T &lower, const T &upper) const { check(l, r); return wm.range_freq(offsets[l], offsets[r], lower, upper); }
	Sum range_sum(int l, int r) const { check(l, r); return wm.range_sum(offsets[l], offsets[r]); }
	Sum range_sum(int l, int r, const T &upper) const { check(l, r); return wm.range_sum(offsets[l], offsets[r], upper); }
	Sum range_sum(int l, int r, const T &lower, const T &upper) const { check(l, r); return wm.range_sum(offsets[l], offsets[r], lower, upper); }
	T kth_smallest(int l, int r, int k) const { check(l, r); return wm.kth_smallest(offsets[l], offsets[r], k); }
	T kth_largest(int l, int r, int k) const { check(l, r); return wm.kth_largest(offsets[l], offsets[r], k); }
	Sum sum_smallest(int l, int r, int k) const { check(l, r); return wm.sum_smallest(offsets[l], offsets[r], k); }
	Sum sum_largest(int l, int r, int k) const { check(l, r); return wm.sum_largest(offsets[l], offsets[r], k); }
	int min_count_smallest(int l, int r, Sum target) const { check(l, r); return wm.min_count_smallest(offsets[l], offsets[r], target); }
	int min_count_largest(int l, int r, Sum target) const { check(l, r); return wm.min_count_largest(offsets[l], offsets[r], target); }
	int max_count_smallest(int l, int r, Sum budget) const { check(l, r); return wm.max_count_smallest(offsets[l], offsets[r], budget); }
	int max_count_largest(int l, int r, Sum budget) const { check(l, r); return wm.max_count_largest(offsets[l], offsets[r], budget); }
};

#endif
Back to top page