Fun With Heaps.

fun_with_heaps.rs

use std::cmp::Ordering;
use std::fmt::Debug;

// min-heap, for now
#[derive(Debug)]
struct Heap<T: Ord> {
    h: Vec<T>,
    /// min heap if true, otherwise max heap
    min: bool,
}

type Ix = usize;

/// Returns the indices of the left and right child, respectively
fn children(ix: Ix) -> (Ix, Ix) {
    (2 * ix + 1, 2 * ix + 2)
}

/// Returns the index of the parent of ix
fn parent(ix: Ix) -> Ix {
    (ix - 1) / 2
}

impl<T: Ord + Debug> Heap<T> {
    pub fn new_min() -> Heap<T> {
        Heap { h: Vec::new(), min: true }
    }
    pub fn new_max() -> Heap<T> {
        Heap { h: Vec::new(), min: false }
    }

    /// Reset the internal vector to v.
    pub fn _set_vec(&mut self, v: Vec<T>) {
        self.h = v;
    }

    /// Swaps the elements at ix1 and ix2.
    fn swap(&mut self, ix1: Ix, ix2: Ix) {
        self.h.swap(ix1, ix2)
    }

    /// Compares two elements.
    /// The rest of the implementation assumes a min-heap; if this heap is configured to be a max
    /// heap, cmp() simply reverses the result.
    fn cmp(&self, ix1: Ix, ix2: Ix) -> Ordering {
        if self.min {
            self.h[ix1].cmp(&self.h[ix2])
        } else {
            self.h[ix1].cmp(&self.h[ix2]).reverse()
        }
    }

    /// Insert an element at the bottom of the heap
    fn insert_bottom(&mut self, elem: T) -> Ix {
        self.h.push(elem);
        self.h.len() - 1
    }

    /// Reorders the element ix so the heap invariant is fulfilled.
    /// Returns the new index.
    fn reorder(&mut self, ix: Ix) -> Ix {
        let mut current_index = ix;
        loop {
            if current_index == 0 {
                break;
            }

            // This is a fixed min-heap right now
            match self.cmp(current_index, parent(current_index)) {
                Ordering::Less => {
                    self.swap(current_index, parent(current_index));
                    current_index = parent(current_index);
                }
                Ordering::Equal => break,
                Ordering::Greater => break,
            }
        }
        current_index
    }

    /// "Bubbles down" the root element.
    /// Needed for removing: First and last element are swapped, last element is removed, first
    /// element is reordered so the heap is satisfied again.
    fn reorder_top(&mut self) -> Ix {
        let mut current_ix = 0;
        loop {
            let (left, right) = children(current_ix);

            if left < self.h.len() && right < self.h.len() {
                if self.cmp(current_ix, left) != Ordering::Less || self.cmp(current_ix, right) != Ordering::Less {
                    match self.cmp(left, right) {
                        Ordering::Less => {
                            self.swap(current_ix, left);
                            current_ix = left;
                        }
                        Ordering::Equal => {
                            self.swap(current_ix, left);
                            current_ix = left;
                        }
                        Ordering::Greater => {
                            self.swap(current_ix, right);
                            current_ix = right;
                        }
                    }
                } else {
                    break
                }
            } else if left < self.h.len() {
                if self.cmp(current_ix, left) == Ordering::Greater {
                    self.swap(current_ix, left);
                    current_ix = left;
                } else {
                    break
                }
            } else if right < self.h.len() {
                if self.cmp(current_ix, right) == Ordering::Greater {
                    self.swap(current_ix, right);
                    current_ix = right;
                } else {
                    break
                }
            } else {
                break
            }
        }
        current_ix
    }

    pub fn insert(&mut self, elem: T) {
        let ix = self.insert_bottom(elem);
        self.reorder(ix);
    }

    pub fn peek(&self) -> Option<T>
        where T: Ord + Debug + Clone
    {
        if self.h.len() < 1 {
            None
        } else {
            Some(self.h[0].clone())
        }
    }

    /// Returns the top-most element
    pub fn pop(&mut self) -> Option<T> {
        if self.h.len() < 1 {
            None
        } else {
            let len = self.h.len();
            self.swap(0, len-1);
            let elem = self.h.remove(len-1);
            self.reorder_top();
            Some(elem)
        }
    }

    pub fn len(&self) -> usize {
        self.h.len()
    }
    
    pub fn dbgprint(&self, ix: Ix) -> String {
        let (left, right) = children(ix);

        if left < self.h.len() && right < self.h.len() {
            format!("N({:?}, {}, {})",
                    self.h[ix],
                    self.dbgprint(left),
                    self.dbgprint(right))
        } else if left < self.h.len() {
            format!("N({:?}, {}, N())", self.h[ix], self.dbgprint(left))
        } else if right < self.h.len() {
            format!("N({:?}, N(), {})", self.h[ix], self.dbgprint(right))
        } else if ix < self.h.len() {
            format!("{:?}", self.h[ix])
        } else {
            "".to_string()
        }
    }
}

struct MedianTracker {
    lower: Heap<u32>,
    upper: Heap<u32>,
}

impl MedianTracker {
    fn new() -> MedianTracker {
        MedianTracker { upper: Heap::new_min(), lower: Heap::new_max() }
    }

    fn rebalance(&mut self) {
        while (self.lower.len() as i64 - self.upper.len() as i64).abs() > 1 {
            if self.lower.len() < self.upper.len() {
                self.lower.insert(self.upper.pop().unwrap());
            } else {
                self.upper.insert(self.lower.pop().unwrap());
            }
        }
    }

    fn insert(&mut self, n: u32) {
        if n <= self.median() {
            self.lower.insert(n);
        } else {
            self.upper.insert(n);
        }
        self.rebalance();
    }

    pub fn median(&self) -> u32 {
        let (len_lower, len_upper) = (self.lower.len(), self.upper.len());
        if len_lower > len_upper {
            self.lower.peek().unwrap()
        } else if len_upper > len_lower {
            self.upper.peek().unwrap()
        } else {
            (self.lower.peek().unwrap_or(0) + self.upper.peek().unwrap_or(0)) / 2
        }
    }
}

fn main() {}

#[cfg(test)]
mod tests {
    use super::Heap;
    use super::MedianTracker;

    #[test]
    fn test_insert() {
        let mut h = Heap::new_min();
        h.insert(10 as u32);
        h.insert(5);
        h.insert(6);
        h.insert(1);
        h.insert(7);
        h.insert(2);
        h.insert(0);
        h.insert(4);


        println!("{}", h.dbgprint(0));
        assert_eq!(h.pop(), Some(0));
        println!("{}", h.dbgprint(0));
        assert_eq!(h.peek(), Some(1));

        while let Some(n) = h.pop() {
            println!("{}", n);
            println!("{}", h.dbgprint(0));
        }
    }
    #[test]
    fn test_insert_max() {
        let mut h = Heap::new_max();
        h.insert(10 as u32);
        h.insert(5);
        h.insert(6);
        h.insert(1);
        h.insert(7);
        h.insert(2);
        h.insert(0);
        h.insert(4);


        println!("{}", h.dbgprint(0));
        assert_eq!(h.pop(), Some(10));
        println!("{}", h.dbgprint(0));
        assert_eq!(h.peek(), Some(7));
        while let Some(n) = h.pop() {
            println!("{}", n);
            println!("{}", h.dbgprint(0));
        }
    }

    #[test]
    fn test_mediantracker() {
        let mut t = MedianTracker::new();

        for i in 0..10000000 {
            t.insert(i);
        }

        assert_eq!(t.median(), 4999999);
    }

    #[test]
    fn bench_heap() {
        let mut h: Heap<u64> = Heap::new_min();
        h._set_vec(Vec::with_capacity(100000000));

        for i in 0..100000000 {
            h.insert(i);
        }

        assert_eq!(h.peek(), Some(0));
    }
}