Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Rust implementation of sorted iterator based on 2 input iterators
use std::iter::{Iterator, Peekable};
use std::collections::BinaryHeap;
use std::cmp::Ordering;
use std::cell::RefCell;
use std::iter::FromIterator;
struct WrappedIterator<T: Ord, P: Iterator<Item = T>>(RefCell<Peekable<P>>);
impl<T: Ord, P: Iterator<Item = T>> PartialEq for WrappedIterator<T, P> {
fn eq(&self, other: &Self) -> bool {
let mut this = self.0.borrow_mut();
let mut other = other.0.borrow_mut();
let this = this.peek();
let other = other.peek();
if this.is_none() && other.is_none() {
return true;
}
if this.is_none() || other.is_none() {
return false;
}
let this = this.unwrap();
let other = other.unwrap();
*this == *other
}
}
impl<T: Ord, P: Iterator<Item = T>> Eq for WrappedIterator<T, P> {}
impl<T: Ord, P: Iterator<Item = T>> PartialOrd for WrappedIterator<T, P> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
let mut this = self.0.borrow_mut();
let mut other = other.0.borrow_mut();
let this = this.peek();
let other = other.peek();
if this.is_none() || other.is_none() {
return None;
}
let this = this.unwrap();
let other = other.unwrap();
match this.partial_cmp(other) {
None => None,
Some(ord) => {
let ord = match ord {
Ordering::Greater => Ordering::Less,
Ordering::Less => Ordering::Greater,
Ordering::Equal => ord,
};
Some(ord)
}
}
}
}
// http://stackoverflow.com/questions/39949939/how-can-i-implement-a-min-heap-of-f64-with-rusts-binaryheap
impl<T: Ord, P: Iterator<Item = T>> Ord for WrappedIterator<T, P> {
fn cmp(&self, other: &Self) -> Ordering {
self.partial_cmp(other).unwrap()
}
}
struct Foo<T: Ord, P: Iterator<Item = T>> {
heap: BinaryHeap<WrappedIterator<T, P>>,
}
impl<T: Ord, P: Iterator<Item = T>> FromIterator<P> for Foo<T, P> {
fn from_iter<I>(iter: I) -> Self
where I: IntoIterator<Item = P>
{
let iter = iter.into_iter().map(|x| WrappedIterator(RefCell::new(x.peekable())));
Foo { heap: BinaryHeap::from_iter(iter) }
}
}
impl<T, P> Iterator for Foo<T, P>
where T: Ord + Clone,
P: Iterator<Item = T>
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
match self.heap.pop() {
None => None,
Some(iter) => {
let next = {
iter.0.borrow_mut().next()
};
let has_next = iter.0.borrow_mut().peek().is_some();
if has_next {
self.heap.push(iter);
}
next
}
}
}
}
// this struct lets you create an iterator from 2 sub iterators, and it will
// automatically return the interleaved, sorted results without much cost
struct OrderedIterator<I: Ord + Clone, P: Iterator<Item = I>> {
first: Peekable<P>,
second: Peekable<P>,
}
impl<T, P> OrderedIterator<T, P>
where T: Ord + Clone,
P: Iterator<Item = T>
{
pub fn new(first: P, second: P) -> Self {
let first_peekable = first.peekable();
let second_peekable = second.peekable();
OrderedIterator {
first: first_peekable,
second: second_peekable,
}
}
}
impl<I, P> Iterator for OrderedIterator<I, P>
where I: Ord + Clone,
P: Iterator<Item = I>
{
type Item = I;
fn next(&mut self) -> Option<Self::Item> {
let first_is_none = self.first.peek().is_none();
let second_is_none = self.second.peek().is_none();
if first_is_none && second_is_none {
return None;
}
if first_is_none {
return self.second.next();
} else if second_is_none {
return self.first.next();
}
if self.first.peek().unwrap() < self.second.peek().unwrap() {
return self.first.next();
} else {
return self.second.next();
}
}
}
fn main() {
let first_vec = vec![1, 5, 7, 11, 13];
let second_vec = vec![2, 3, 4, 6, 8, 9, 10, 12];
let vec_iters = vec![first_vec.clone().into_iter(), second_vec.clone().into_iter()];
for item in OrderedIterator::new(first_vec.into_iter(), second_vec.into_iter()) {
println!("{}", item);
}
for item in Foo::from_iter(vec_iters) {
println!("{}", item);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment