集合のn番目を取得する方法

[tipstar0125]

RustのBTreeSetは、lower_boundやupper_boundはrangeメソッドで代用可能だが、i-th番目の情報を高速で取得することができない。 BITを使用すると、O(N)の計算量をO(logN)に落とすことができる。

• addメソッドで、集合への追加と集合からの削除ができる
• lower_boundメソッドで、i-th番目の値を取得できる(1-indexed!)

let mut set = BinaryIndexedTree::new(10);
// {1, 3, 5, 7, 9}
for i in 0..10 {
    if i % 2 == 1 {
        set.add(i, 1);
    }
}
// 1-indexed!!
println!("{}", set.lower_bound(0)); // 0(Do not use!)
println!("{}", set.lower_bound(1)); // 1
println!("{}", set.lower_bound(2)); // 3
println!("{}", set.lower_bound(3)); // 5
println!("{}", set.lower_bound(4)); // 7
println!("{}", set.lower_bound(5)); // 9
println!("{}", set.lower_bound(6)); // 10(size)

set.add(5, -1);
set.add(7, -1);
// {1, 3, 9}
println!("{}", set.lower_bound(1)); // 1
println!("{}", set.lower_bound(2)); // 3
println!("{}", set.lower_bound(3)); // 9

[tipstar0125]

BITライブラリ

#[derive(Debug, Clone)]
struct BinaryIndexedTree {
    size: usize,
    data: Vec<isize>,
}

impl BinaryIndexedTree {
    fn new(n: usize) -> Self {
        BinaryIndexedTree {
            size: n,
            data: vec![0; n],
        }
    }
    fn lsb(&self, i: usize) -> usize {
        i & i.wrapping_neg()
    }
    fn build(&mut self, v: &[isize]) {
        assert_eq!(self.size, v.len(), "size not correct!");
        self.data = v.to_vec();
        for i in 1..=self.size {
            let lsb = self.lsb(i);
            if i + lsb <= self.size {
                self.data[i + lsb - 1] += self.data[i - 1];
            }
        }
    }
    fn push(&mut self, mut x: isize) {
        self.size += 1;
        let mut d = 1;
        let k = self.lsb(self.size);
        while d != k {
            x += self.data[self.size - d - 1];
            d <<= 1;
        }
        self.data.push(x);
    }
    fn add(&mut self, i: usize, x: isize) {
        let mut idx = i + 1;
        while idx <= self.size {
            self.data[idx - 1] += x;
            idx += self.lsb(idx);
        }
    }
    //  [0, r)
    fn sum(&self, i: usize) -> isize {
        let mut ret = 0;
        let mut idx = i;
        while idx > 0 {
            ret += self.data[idx - 1];
            idx -= self.lsb(idx);
        }
        ret
    }
    // [l, r)
    fn range_sum(&self, l: usize, r: usize) -> isize {
        self.sum(r) - self.sum(l)
    }
    fn lower_bound(&self, x: isize) -> usize {
        let mut i = 0;
        let mut k = 1;
        let mut x = x;
        while k <= self.size {
            k <<= 1;
        }
        while k > 0 {
            if i + k <= self.size && self.data[i + k - 1] < x {
                x -= self.data[i + k - 1];
                i += k;
            }
            k >>= 1;
        }
        if x > 0 {
            i
        } else {
            0
        }
    }
    fn upper_bound(&self, x: isize) -> usize {
        let mut i = 0;
        let mut k = 1;
        let mut x = x;
        while k <= self.size {
            k <<= 1;
        }
        while k > 0 {
            if i + k <= self.size && self.data[i + k - 1] <= x {
                x -= self.data[i + k - 1];
                i += k;
            }
            k >>= 1;
        }
        if i < self.size {
            i
        } else {
            self.size
        }
    }
}

[tipstar0125]

使用例：https://yukicoder.me/submissions/871275