Skip to content

Instantly share code, notes, and snippets.

@ssrlive
Last active March 6, 2023 15:48
Show Gist options
  • Save ssrlive/174005e0ead1e69d0302b8f5860b69ae to your computer and use it in GitHub Desktop.
Save ssrlive/174005e0ead1e69d0302b8f5860b69ae to your computer and use it in GitHub Desktop.
lower_bound and upper_bound in rust.
use std::{cmp::Ordering, result::Result};
/// find first index where arr[idx] >= v; assume arr is sorted
pub trait LowerBound {
type Item;
fn lower_bound(&self, x: &Self::Item) -> Result<usize, usize>;
}
/// find first index where arr[idx] > v; assume arr is sorted
trait UpperBound {
type Item;
fn upper_bound(&self, x: &Self::Item) -> Result<usize, usize>;
}
impl<T: Ord> LowerBound for [T] {
type Item = T;
fn lower_bound(&self, x: &Self::Item) -> Result<usize, usize> {
let mut left = 0;
let len = self.len();
let mut right = len;
while left < right {
let mid = left + (right - left) / 2;
match self[mid].cmp(x) {
Ordering::Less => left = mid + 1,
_ => right = mid,
}
}
assert_eq!(left, right);
if left == len {
Err(left)
} else {
Ok(left)
}
}
}
impl<T: Ord> UpperBound for [T] {
type Item = T;
fn upper_bound(&self, x: &Self::Item) -> Result<usize, usize> {
let mut left = 0;
let len = self.len();
let mut right = len;
while left < right {
let mid = left + (right - left) / 2;
match self[mid].cmp(x) {
Ordering::Greater => right = mid,
_ => left = mid + 1,
}
}
assert_eq!(left, right);
if left == len {
Err(left)
} else {
Ok(left)
}
}
}
impl<T: Ord> LowerBound for Vec<T> {
type Item = T;
fn lower_bound(&self, x: &Self::Item) -> Result<usize, usize> {
self.as_slice().lower_bound(x)
}
}
impl<T: Ord> UpperBound for Vec<T> {
type Item = T;
fn upper_bound(&self, x: &Self::Item) -> Result<usize, usize> {
self.as_slice().upper_bound(x)
}
}
use std::collections::*;
impl<T: Ord> LowerBound for VecDeque<T> {
type Item = T;
fn lower_bound(&self, x: &Self::Item) -> Result<usize, usize> {
self.as_slices().0.lower_bound(x)
}
}
impl<T: Ord> UpperBound for VecDeque<T> {
type Item = T;
fn upper_bound(&self, x: &Self::Item) -> Result<usize, usize> {
self.as_slices().0.upper_bound(x)
}
}
impl<T: Ord> LowerBound for BinaryHeap<T> {
type Item = T;
fn lower_bound(&self, x: &Self::Item) -> Result<usize, usize> {
self.iter().position(|y| y >= x).ok_or(self.len())
}
}
impl<T: Ord> UpperBound for BinaryHeap<T> {
type Item = T;
fn upper_bound(&self, x: &Self::Item) -> Result<usize, usize> {
self.iter().position(|y| y > x).ok_or(self.len())
}
}
#[test]
fn test_lower_bound() {
let v = Vec::<i32>::new();
assert_eq!(v.lower_bound(&0), Err(0));
assert_eq!(v.lower_bound(&1), Err(0));
let v = vec![1, 2, 4, 5, 5, 6, 6];
assert_eq!(v.lower_bound(&0), Ok(0));
assert_eq!(v.lower_bound(&1), Ok(0));
assert_eq!(v.lower_bound(&2), Ok(1));
assert_eq!(v.lower_bound(&3), Ok(2));
assert_eq!(v.lower_bound(&4), Ok(2));
assert_eq!(v.lower_bound(&5), Ok(3));
assert_eq!(v.lower_bound(&6), Ok(5));
assert_eq!(v.lower_bound(&7), Err(7));
assert_eq!(v.lower_bound(&8), Err(7));
assert_eq!(v.lower_bound(&9), Err(7));
}
#[test]
fn test_upper_bound() {
let v = Vec::<i32>::new();
assert_eq!(v.upper_bound(&0), Err(0));
assert_eq!(v.upper_bound(&1), Err(0));
let v = vec![-1, 1, 2, 4, 5, 5, 6, 6];
assert_eq!(v.upper_bound(&0), Ok(1));
assert_eq!(v.upper_bound(&1), Ok(2));
assert_eq!(v.upper_bound(&2), Ok(3));
assert_eq!(v.upper_bound(&3), Ok(3));
assert_eq!(v.upper_bound(&4), Ok(4));
assert_eq!(v.upper_bound(&5), Ok(6));
assert_eq!(v.upper_bound(&6), Err(8));
assert_eq!(v.upper_bound(&7), Err(8));
assert_eq!(v.upper_bound(&8), Err(8));
assert_eq!(v.upper_bound(&9), Err(8));
}
use std::collections::VecDeque;
pub fn lower_bound<T: Ord>(deque: &VecDeque<T>, x: &T) -> Result<usize, usize>
{
let len = deque.len();
let count = deque.iter().filter(|&e| e <= x).count();
let idx = deque.binary_search_by(|e| e.cmp(x));
match idx {
Ok(idx) => Ok(idx),
Err(idx) => {
if idx == len {
Err(count)
} else {
Ok(idx)
}
}
}
}
pub fn upper_bound<T: Ord>(deque: &VecDeque<T>, x: &T) -> Result<usize, usize>
{
let len = deque.len();
let position = deque.iter().position(|e| e > x);
let idx = deque.binary_search_by(|e| e.cmp(x));
match position {
Some(index) => Ok(index),
None => {
match idx {
Ok(_) => Err(len),
Err(idx) => {
if idx == len {
Err(len)
} else {
Ok(idx)
}
}
}
}
}
}
#[test]
fn test() {
let nums = vec![1, 2, 4, 5, 5, 6, 6];
let q = nums.iter().map(|&x| x as i64).collect::<VecDeque<i64>>();
assert_eq!(lower_bound(&q, &0), Ok(0));
assert_eq!(lower_bound(&q, &1), Ok(0));
assert_eq!(lower_bound(&q, &2), Ok(1));
assert_eq!(lower_bound(&q, &3), Ok(2));
assert_eq!(lower_bound(&q, &4), Ok(2));
assert_eq!(lower_bound(&q, &5), Ok(3));
assert_eq!(lower_bound(&q, &6), Ok(5));
assert_eq!(lower_bound(&q, &7), Err(7));
assert_eq!(lower_bound(&q, &8), Err(7));
assert_eq!(lower_bound(&q, &9), Err(7));
let nums = vec![-1, 1, 2, 4, 5, 5, 6, 6];
let q = nums.iter().map(|&x| x as i64).collect::<VecDeque<i64>>();
assert_eq!(upper_bound(&q, &0), Ok(1));
assert_eq!(upper_bound(&q, &1), Ok(2));
assert_eq!(upper_bound(&q, &2), Ok(3));
assert_eq!(upper_bound(&q, &3), Ok(3));
assert_eq!(upper_bound(&q, &4), Ok(4));
assert_eq!(upper_bound(&q, &5), Ok(6));
assert_eq!(upper_bound(&q, &6), Err(8));
assert_eq!(upper_bound(&q, &7), Err(8));
assert_eq!(upper_bound(&q, &8), Err(8));
assert_eq!(upper_bound(&q, &9), Err(8));
}
@ssrlive
Copy link
Author

ssrlive commented Mar 3, 2023

This code can do all of things above:

    fn lower_bound<K: Ord, V>(btm: &BTreeMap<K, V>, x: &K) -> Result<usize, usize> {
        // btm.iter().position(|(k, _)| *k >= *x).ok_or(btm.len())
        use std::cmp::Ordering::Less;
        btm.iter().position(|(k, _)| k.cmp(x) != Less).ok_or(btm.len())
    }


fn main() {
    use itertools::partition;
    let mut data = [7, 1, 1, 7, 1, 1, 7];
    let split_index = partition(&mut data, |elt| *elt >= 3);
    assert_eq!(data, [7, 7, 7, 1, 1, 1, 1]);
    assert_eq!(split_index, 3);

    let a = [1, 2, 3];
    let (even, odd): (Vec<i32>, Vec<i32>) = a.into_iter().partition(|&n| n % 2 == 0);
    assert_eq!(even, vec![2]);
    assert_eq!(odd, vec![1, 3]);
}

@ssrlive
Copy link
Author

ssrlive commented Mar 6, 2023

use std::cmp::Ordering;
fn lower_bound_by<T, F>(arr: &[T], f: F) -> Result<usize, usize>
where
    T: Ord,
    F: Fn(&T) -> Ordering,
{
    arr.iter().position(|y| f(y) != Ordering::Less).ok_or(arr.len())
}

fn lower_bound<T: Ord>(arr: &[T], x: &T) -> Result<usize, usize> {
    lower_bound_by(arr, |y| y.cmp(x))
}

fn upper_bound_by<T, F>(arr: &[T], f: F) -> Result<usize, usize>
where
    T: Ord,
    F: Fn(&T) -> Ordering,
{
    arr.iter().position(|y| f(y) == Ordering::Greater).ok_or(arr.len())
}

fn upper_bound<T: Ord>(arr: &[T], x: &T) -> Result<usize, usize> {
    upper_bound_by(arr, |y| y.cmp(x))
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment