-
-
Save dashed/2593b3a6e321dde4eaf27b814f6534bc to your computer and use it in GitHub Desktop.
Rust implementation of sorted iterator based on 2 input iterators
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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