概要
BTreeSet は「集合に入った要素のうち x 以上 / 以下を取りたい」などの操作が O(log N) ででき、競技プログラミングではとても便利な標準ライブラリです。
一方で BTreeSet 単体では、次のような操作ができません。
- 集合の昇順n番目の要素を取得する
- 区間 [L, R)に含まれる要素数を取得する
たまにこれらが欲しくなる問題があるので、BTreeSetとFenwickTreeを組み合わせたOrderedSetをライブラリ化しました。
実装
早速ですが実装は以下です。
この実装は オフラインを想定しています。
あらかじめ登場する値を from(&Vec) に渡して座標圧縮しておき、その後 insert/remove などを行います。
FenwickTreeのlower_boundを実装しているため、get_nthはO(log N)で取れます。
pub mod ordered_set {
use num_traits::Bounded;
use std::{
collections::{BTreeSet, HashMap},
hash::Hash,
};
use std::ops::{Bound, RangeBounds};
#[derive(Clone)]
pub struct FenwickTree<T> {
n: usize,
ary: Vec<T>,
e: T,
}
impl<T: Clone + std::ops::AddAssign<T> + PartialOrd> FenwickTree<T> {
pub fn new(n: usize, e: T) -> Self {
FenwickTree {
n,
ary: vec![e.clone(); n],
e,
}
}
pub fn accum(&self, mut idx: usize) -> T {
let mut sum = self.e.clone();
while idx > 0 {
sum += self.ary[idx - 1].clone();
idx &= idx - 1;
}
sum
}
/// performs data[idx] += val;
pub fn add<U: Clone>(&mut self, mut idx: usize, val: U)
where
T: std::ops::AddAssign<U>,
{
let n = self.n;
idx += 1;
while idx <= n {
self.ary[idx - 1] += val.clone();
idx += idx & idx.wrapping_neg();
}
}
/// Returns data[l] + ... + data[r - 1].
pub fn sum<R>(&self, range: R) -> T
where
T: std::ops::Sub<Output = T>,
R: RangeBounds<usize>,
{
let r = match range.end_bound() {
Bound::Included(r) => r + 1,
Bound::Excluded(r) => *r,
Bound::Unbounded => self.n,
};
let l = match range.start_bound() {
Bound::Included(l) => *l,
Bound::Excluded(l) => l + 1,
Bound::Unbounded => return self.accum(r),
};
self.accum(r) - self.accum(l)
}
/// prefix sum が k 以上になる最小の idx (0-indexed)
/// 値が全て非負であること仮定。
pub fn lower_bound(&self, k: T) -> Option<usize> {
// k <= 0 相当なら先頭（eが0のとき）。一般型では厳密判定しづらいので e と比較。
if k <= self.e {
return Some(0);
}
if self.accum(self.n) < k {
return None;
}
let mut idx = 0usize;
let mut cur = self.e.clone();
let mut bit = 1usize;
while (bit << 1) <= self.n {
bit <<= 1;
}
while bit > 0 {
let next = idx + bit;
if next <= self.n {
let mut cand = cur.clone();
cand += self.ary[next - 1].clone();
if cand < k {
cur = cand;
idx = next;
}
}
bit >>= 1;
}
Some(idx) // 0-indexed の data の位置
}
/// prefix sum が k より大きくなる最小 idx (0-indexed)
/// 値が全て非負であること仮定。
pub fn upper_bound(&self, k: T) -> Option<usize> {
if self.accum(self.n) <= k {
return None;
}
let mut idx = 0usize;
let mut cur = self.e.clone();
let mut bit = 1usize;
while (bit << 1) <= self.n {
bit <<= 1;
}
while bit > 0 {
let next = idx + bit;
if next <= self.n {
let mut cand = cur.clone();
cand += self.ary[next - 1].clone();
if cand <= k {
cur = cand;
idx = next;
}
}
bit >>= 1;
}
Some(idx)
}
}
pub struct OrderedSet<T: Ord + Hash + Copy + Bounded> {
id_map: HashMap<T, usize>,
inv: Vec<T>,
id_max: usize,
fen: FenwickTree<isize>,
set: BTreeSet<T>,
}
impl<T: Ord + Hash + Copy + Bounded> OrderedSet<T> {
/// オフラインでの利用を想定
pub fn from(v: &Vec<T>) -> Self {
let mut v = v.to_owned();
v.sort();
v.dedup();
let mut id = 0;
let mut id_map = HashMap::new();
for &num in v.iter() {
id_map.insert(num, id);
id += 1usize;
}
let mut inv = vec![T::max_value(); id];
let fen = FenwickTree::new(id + 5, 0);
for &num in v.iter() {
let &id = id_map.get(&num).unwrap();
inv[id] = num;
}
Self {
id_map,
inv,
id_max: id,
fen,
set: BTreeSet::new(),
}
}
pub fn insert(&mut self, num: T) -> bool {
let id = self.get_id(num);
if !self.set.insert(num) {
return false;
}
self.fen.add(id, 1);
true
}
pub fn remove(&mut self, num: T) -> bool {
let id = self.get_id(num);
if !self.set.remove(&num) {
return false;
}
self.fen.add(id, -1);
true
}
pub fn range_count<R>(&self, range: R) -> usize
where
R: RangeBounds<T>,
{
let set = self.set();
let low: Option<T> = match range.start_bound() {
Bound::Included(&l) => set.range(l..).next().copied(),
Bound::Excluded(&l) => set
.range((Bound::Excluded(l), Bound::Unbounded))
.next()
.copied(),
Bound::Unbounded => set.iter().next().copied(),
};
let high: Option<T> = match range.end_bound() {
Bound::Included(&r) => set.range(..=r).next_back().copied(),
Bound::Excluded(&r) => set.range(..r).next_back().copied(),
Bound::Unbounded => set.iter().next_back().copied(),
};
let (low, high) = match (low, high) {
(Some(a), Some(b)) => (a, b),
_ => return 0,
};
assert!(low <= high);
let l = self.get_id(low);
let r = self.get_id(high);
self.fen.sum(l..=r) as usize
}
/// 昇順n番目の要素を取得する。0-indexed
pub fn get_nth(&self, n: usize) -> Option<T> {
let c = n + 1;
let id = self.fen.lower_bound(c as isize)?;
Some(self.inv[id])
}
pub fn set(&self) -> &BTreeSet<T> {
&self.set
}
pub fn get_id(&self, num: T) -> usize {
let &id = self.id_map.get(&num).unwrap();
id
}
pub fn max_id(&self) -> usize {
self.id_max
}
}
}
各APIの概要と計算量
本実装ではオフラインで使用されることが想定されており、
あらかじめinsertされる値をfromでbuildしてからの利用が想定されています。
そのため、実際の問題で利用する場合はクエリの先読みなどが必要になる認識です。
以下ではsetの要素数 := Nとして計算量を表現しています。
| API | 概要 | 計算量 |
|---|---|---|
| insert | setへのinsert | O(log N) |
| remove | setからのremove | O(log N) |
| range_count | 引数で指定したrangeの中に存在する要素数を返す | O(log N) |
| get_nth | 昇順N番目の要素の取得 | O(log N) |
| set | 内部で扱っているBTreeSetのGetter | O(1) |
実装概要
BTreeSet だけではget_nthや range_countが直接できないので、内部で FenwickTree（BIT） を持って解決しています。
- BTreeSet：値の存在管理、range による境界探索
- FenwickTree：座圧した index 上で存在する（=1）/しない（=0）を管理
verify
use proconio::input;
pub trait CollectVec: Iterator + Sized {
fn collect_vec(self) -> Vec<Self::Item>;
}
impl<I: Iterator> CollectVec for I {
fn collect_vec(self) -> Vec<Self::Item> {
self.collect()
}
}
pub trait Joiner {
fn join2(self, sep: &str) -> String;
}
impl<T: ToString, I: IntoIterator<Item = T>> Joiner for I {
fn join2(self, sep: &str) -> String {
self.into_iter()
.map(|x| x.to_string())
.collect::<Vec<_>>()
.join(sep)
}
}
use ordered_set::OrderedSet;
pub mod ordered_set {
use num_traits::Bounded;
use std::{
collections::{BTreeSet, HashMap},
hash::Hash,
};
use std::ops::{Bound, RangeBounds};
#[derive(Clone)]
pub struct FenwickTree<T> {
n: usize,
ary: Vec<T>,
e: T,
}
impl<T: Clone + std::ops::AddAssign<T> + PartialOrd> FenwickTree<T> {
pub fn new(n: usize, e: T) -> Self {
FenwickTree {
n,
ary: vec![e.clone(); n],
e,
}
}
pub fn accum(&self, mut idx: usize) -> T {
let mut sum = self.e.clone();
while idx > 0 {
sum += self.ary[idx - 1].clone();
idx &= idx - 1;
}
sum
}
/// performs data[idx] += val;
pub fn add<U: Clone>(&mut self, mut idx: usize, val: U)
where
T: std::ops::AddAssign<U>,
{
let n = self.n;
idx += 1;
while idx <= n {
self.ary[idx - 1] += val.clone();
idx += idx & idx.wrapping_neg();
}
}
/// Returns data[l] + ... + data[r - 1].
pub fn sum<R>(&self, range: R) -> T
where
T: std::ops::Sub<Output = T>,
R: RangeBounds<usize>,
{
let r = match range.end_bound() {
Bound::Included(r) => r + 1,
Bound::Excluded(r) => *r,
Bound::Unbounded => self.n,
};
let l = match range.start_bound() {
Bound::Included(l) => *l,
Bound::Excluded(l) => l + 1,
Bound::Unbounded => return self.accum(r),
};
self.accum(r) - self.accum(l)
}
/// prefix sum が k 以上になる最小の idx (0-indexed)
/// 値が全て非負であること仮定。
pub fn lower_bound(&self, k: T) -> Option<usize> {
// k <= 0 相当なら先頭（eが0のとき）。一般型では厳密判定しづらいので e と比較。
if k <= self.e {
return Some(0);
}
if self.accum(self.n) < k {
return None;
}
let mut idx = 0usize;
let mut cur = self.e.clone();
let mut bit = 1usize;
while (bit << 1) <= self.n {
bit <<= 1;
}
while bit > 0 {
let next = idx + bit;
if next <= self.n {
let mut cand = cur.clone();
cand += self.ary[next - 1].clone();
if cand < k {
cur = cand;
idx = next;
}
}
bit >>= 1;
}
Some(idx) // 0-indexed の data の位置
}
/// prefix sum が k より大きくなる最小 idx (0-indexed)
/// 値が全て非負であること仮定。
pub fn upper_bound(&self, k: T) -> Option<usize> {
if self.accum(self.n) <= k {
return None;
}
let mut idx = 0usize;
let mut cur = self.e.clone();
let mut bit = 1usize;
while (bit << 1) <= self.n {
bit <<= 1;
}
while bit > 0 {
let next = idx + bit;
if next <= self.n {
let mut cand = cur.clone();
cand += self.ary[next - 1].clone();
if cand <= k {
cur = cand;
idx = next;
}
}
bit >>= 1;
}
Some(idx)
}
}
pub struct OrderedSet<T: Ord + Hash + Copy + Bounded> {
id_map: HashMap<T, usize>,
inv: Vec<T>,
id_max: usize,
fen: FenwickTree<isize>,
set: BTreeSet<T>,
}
impl<T: Ord + Hash + Copy + Bounded> OrderedSet<T> {
/// オフラインでの利用を想定
pub fn from(v: &Vec<T>) -> Self {
let mut v = v.to_owned();
v.sort();
v.dedup();
let mut id = 0;
let mut id_map = HashMap::new();
for &num in v.iter() {
id_map.insert(num, id);
id += 1usize;
}
let mut inv = vec![T::max_value(); id];
let fen = FenwickTree::new(id + 5, 0);
for &num in v.iter() {
let &id = id_map.get(&num).unwrap();
inv[id] = num;
}
Self {
id_map,
inv,
id_max: id,
fen,
set: BTreeSet::new(),
}
}
pub fn insert(&mut self, num: T) -> bool {
let id = self.get_id(num);
if !self.set.insert(num) {
return false;
}
self.fen.add(id, 1);
true
}
pub fn remove(&mut self, num: T) -> bool {
let id = self.get_id(num);
if !self.set.remove(&num) {
return false;
}
self.fen.add(id, -1);
true
}
pub fn range_count<R>(&self, range: R) -> usize
where
R: RangeBounds<T>,
{
let set = self.set();
let low: Option<T> = match range.start_bound() {
Bound::Included(&l) => set.range(l..).next().copied(),
Bound::Excluded(&l) => set
.range((Bound::Excluded(l), Bound::Unbounded))
.next()
.copied(),
Bound::Unbounded => set.iter().next().copied(),
};
let high: Option<T> = match range.end_bound() {
Bound::Included(&r) => set.range(..=r).next_back().copied(),
Bound::Excluded(&r) => set.range(..r).next_back().copied(),
Bound::Unbounded => set.iter().next_back().copied(),
};
let (low, high) = match (low, high) {
(Some(a), Some(b)) => (a, b),
_ => return 0,
};
assert!(low <= high);
let l = self.get_id(low);
let r = self.get_id(high);
self.fen.sum(l..=r) as usize
}
/// 昇順n番目の要素を取得する。0-indexed
pub fn get_nth(&self, n: usize) -> Option<T> {
let c = n + 1;
let id = self.fen.lower_bound(c as isize)?;
Some(self.inv[id])
}
pub fn set(&self) -> &BTreeSet<T> {
&self.set
}
pub fn get_id(&self, num: T) -> usize {
let &id = self.id_map.get(&num).unwrap();
id
}
pub fn max_id(&self) -> usize {
self.id_max
}
}
}
fn main() {
input! {
n: usize,
q: usize,
a: [usize; n],
tx: [(usize, usize); q]
}
let mut xs = tx.iter().map(|&e| e.1).collect_vec();
xs.extend(a.to_owned());
let mut set = OrderedSet::from(&xs);
for &num in a.iter() {
set.insert(num);
}
for &(t, x) in tx.iter() {
let s = set.set();
if t == 0 {
set.insert(x);
} else if t == 1 {
set.remove(x);
} else if t == 2 {
let idx = set.get_nth(x - 1);
if idx.is_none() {
println!("-1");
continue;
}
let idx = idx.unwrap();
println!("{}", idx);
} else if t == 3 {
let ans = set.range_count(..=x);
println!("{}", ans);
} else if t == 4 {
let ll = s.range(..=x).next_back();
if ll.is_none() {
println!("-1");
continue;
}
let ll = ll.unwrap();
println!("{}", ll);
} else {
let rr = s.range(x..).next();
if rr.is_none() {
println!("-1");
continue;
}
let rr = rr.unwrap();
println!("{}", rr);
}
}
}
おわりに
できればオンラインで扱いたいです。
もし何か方法があればコメントいただければ幸いです。