algorithms_fourth 0.1.10

//!
//! # 使优先级队列用
//!
//! ```
//! use algorithms_fourth::priority_queue::PQ;
//!    let mut pq = PQ::new_min_pq();
//!    pq.push(2);
//!    pq.push(1);
//!    pq.push(5);
//!    pq.push(3);
//!    pq.push(4);
//!    pq.push(6);
//!    assert_eq!(pq.peek(), Some(&1));
//!    assert_eq!(pq.peek(), Some(&1));
//!    assert_eq!(pq.pop(), Some(1));
//!    assert_eq!(pq.peek(), Some(&2));
//!    assert_eq!(pq.pop(), Some(2));
//!    assert_eq!(pq.pop(), Some(3));
//!    assert_eq!(pq.pop(), Some(4));
//!    assert_eq!(pq.pop(), Some(5));
//!    assert_eq!(pq.pop(), Some(6));
//!    assert_eq!(pq.pop(), None);
//! ```
//!
//! # 索引优先级队列使用
//!
//! ```
//! use algorithms_fourth::priority_queue::IndexPQ;
//!
//!    let mut pq = IndexPQ::new_min_pq(6);
//!    pq.push(0, 2);
//!    pq.push(1, 1);
//!    pq.push(2, 5);
//!    pq.push(3, 3);
//!    pq.push(4, 4);
//!    pq.push(5, 6);
//!    assert_eq!(pq.peek(), Some(( 1,&1 )));
//!    assert_eq!(pq.pop(), Some(( 1,1 )));
//!    pq.change(0, 7);
//!    assert_eq!(pq.pop(), Some((3,3)));
//!    assert_eq!(pq.pop(), Some(( 4,4 )));
//!    assert_eq!(pq.pop(), Some((2,5)));
//!    assert_eq!(pq.pop(), Some((5,6)));
//!    assert_eq!(pq.pop(), Some((0,7)));
//!    assert_eq!(pq.pop(), None);
//!
//!    let mut pq = IndexPQ::new_max_pq(6);
//!    pq.push(0, 2);
//!    pq.push(1, 1);
//!    pq.push(2, 5);
//!    pq.push(3, 3);
//!    pq.push(4, 4);
//!    pq.push(5, 6);
//!    assert_eq!(pq.peek(), Some((5, &6 )));
//!    assert_eq!(pq.pop(), Some((5,6)));
//!    pq.change(4, 7);
//!    assert_eq!(pq.pop(), Some((4,7)));
//!    assert_eq!(pq.pop(), Some((2,5)));
//!    assert_eq!(pq.peek(), Some((3,&3)));
//!    assert_eq!(pq.pop(), Some((3,3)));
//!    assert_eq!(pq.pop(), Some((0,2)));
//!    assert_eq!(pq.pop(), Some((1,1)));
//!    
//! ```
//!  

pub struct PQ<T> {
    data: Vec<T>,
    need_swap: fn(&T, &T) -> bool,
}
impl<T: PartialOrd> PQ<T> {
    pub fn new_min_pq() -> Self {
        PQ {
            data: Vec::new(),
            need_swap: |a, b| a > b,
        }
    }
    pub fn new_max_pq() -> Self {
        PQ {
            data: Vec::new(),
            need_swap: |a, b| a < b,
        }
    }
    pub fn push(&mut self, val: T) {
        let pos = self.data.len();
        self.data.push(val);
        self.swim(pos);
    }
    fn get_parent_pos(pos: usize) -> Option<usize> {
        if 0 == pos {
            None
        } else {
            let parent = (pos + 1) / 2 - 1;
            Some(parent)
        }
    }
    fn get_child_pos(pos: usize, len: usize) -> Option<(usize, Option<usize>)> {
        let left = (pos + 1) * 2 - 1;
        if left >= len {
            None
        } else if left == len - 1 {
            Some((left, None))
        } else {
            Some((left, Some(left + 1)))
        }
    }
    fn sink(&mut self, pos: usize) {
        let mut i = pos;
        let len = self.data.len();
        let data = &mut self.data;
        while let Some(child_pos) = Self::get_child_pos(i, len) {
            let (l, r) = child_pos;
            let mut target = i;
            if (self.need_swap)(&data[target], &data[l]) {
                target = l;
            }
            if let Some(r) = r {
                if (self.need_swap)(&data[target], &data[r]) {
                    target = r;
                }
            }
            if target != i {
                data.swap(target, i);
                i = target;
            } else {
                break;
            }
        }
    }
    fn swim(&mut self, pos: usize) {
        let mut i = pos;
        while let Some(p) = Self::get_parent_pos(i) {
            if (self.need_swap)(&self.data[p], &self.data[i]) {
                self.data.swap(p, i);
                i = p;
            } else {
                break;
            }
        }
    }
    pub fn peek(&self) -> Option<&T> {
        self.data.get(0)
    }
    pub fn size(&self) -> usize {
        self.data.len()
    }
    pub fn is_empty(&self) -> bool {
        self.data.get(0).is_none()
    }
    pub fn pop(&mut self) -> Option<T> {
        if self.data.is_empty() {
            None
        } else {
            let ans = Some(self.data.swap_remove(0));
            self.sink(0);
            ans
        }
    }
}

/// 索引优先级队列
pub struct IndexPQ<T> {
    /// 有优先级之分的元素
    keys: Vec<Option<T>>,
    /// 索引二叉堆
    data: Vec<usize>,
    /// 逆序。keys中第i个元素在data中的对应位置
    /// qp\[data\[i]] = data\[qp\[i]] = i
    qp: Vec<usize>,
    /// 当后者是否优于前者
    need_swap: fn(&T, &T) -> bool,
}
impl<T: PartialOrd> IndexPQ<T> {
    pub fn new_min_pq(count: usize) -> Self {
        let mut keys = Vec::with_capacity(count);
        let qp = vec![0; count];
        for _ in 0..count {
            keys.push(None);
        }
        IndexPQ {
            keys,
            qp,
            data: Vec::new(),
            need_swap: |a, b| a > b,
        }
    }
    pub fn new_max_pq(count: usize) -> Self {
        let mut keys = Vec::with_capacity(count);
        let qp = vec![0; count];
        for _ in 0..count {
            keys.push(None);
        }
        IndexPQ {
            keys,
            data: Vec::new(),
            qp,
            need_swap: |a, b| a < b,
        }
    }
    pub fn push(&mut self, k: usize, val: T) {
        assert!(self.keys.len() > k);
        assert!(self.keys[k].is_none(), "修改数据请用change()");
        self.keys[k] = Some(val);
        self.data.push(k);
        let pos = self.data.len() - 1;
        self.qp[k] = pos;
        self.swim(pos);
    }
    fn get_parent_pos(pos: usize) -> Option<usize> {
        if 0 == pos {
            None
        } else {
            let parent = (pos + 1) / 2 - 1;
            Some(parent)
        }
    }
    fn get_child_pos(pos: usize, len: usize) -> Option<(usize, Option<usize>)> {
        let left = (pos + 1) * 2 - 1;
        if left >= len {
            None
        } else if left == len - 1 {
            Some((left, None))
        } else {
            Some((left, Some(left + 1)))
        }
    }
    fn sink(&mut self, pos: usize) {
        let mut i = pos;
        let len = self.data.len();
        while let Some(child_pos) = Self::get_child_pos(i, len) {
            let (l, r) = child_pos;
            let mut target = i;
            let data = &self.data;
            if (self.need_swap)(self.get_key(data[target]), self.get_key(data[l])) {
                target = l;
            }
            if let Some(r) = r {
                if (self.need_swap)(self.get_key(data[target]), self.get_key(data[r])) {
                    target = r;
                }
            }
            if target != i {
                self.swap(target, i);
                i = target;
            } else {
                break;
            }
        }
    }
    fn get_key(&self, k: usize) -> &T {
        self.keys[k].as_ref().unwrap()
    }
    fn swim(&mut self, pos: usize) {
        let mut i = pos;
        while let Some(p) = Self::get_parent_pos(i) {
            if (self.need_swap)(self.get_key(self.data[p]), self.get_key(self.data[i])) {
                self.swap(p, i);
                i = p;
            } else {
                break;
            }
        }
    }
    fn swap(&mut self, a: usize, b: usize) {
        self.data.swap(a, b);
        self.qp.swap(self.data[b], self.data[a])
    }
    pub fn peek(&self) -> Option<(usize, &T)> {
        self.data.first().map(|i| (*i, self.get_key(*i)))
    }
    pub fn pop(&mut self) -> Option<(usize, T)> {
        if self.data.is_empty() {
            None
        } else {
            self.swap(0, self.data.len() - 1);
            let k = self.data.pop().unwrap();
            let key = self.keys[k].take().unwrap();
            // 修正优先级
            self.sink(0);
            Some((k, key))
        }
    }
    pub fn change(&mut self, k: usize, key: T) {
        self.keys[k] = Some(key);
        self.swim(self.qp[k]);
        self.sink(self.qp[k]);
    }
    pub fn contains(&self, k: usize) -> bool {
        self.keys[k].is_some()
    }
}

#[cfg(test)]
mod test {
    use crate::priority_queue::IndexPQ;

    use super::PQ;

    #[test]
    fn test_min_pq() {
        let mut pq = PQ::new_min_pq();
        pq.push(2);
        pq.push(1);
        pq.push(5);
        pq.push(3);
        pq.push(4);
        pq.push(6);
        assert_eq!(pq.peek(), Some(&1));
        assert_eq!(pq.pop(), Some(1));
        assert_eq!(pq.pop(), Some(2));
        assert_eq!(pq.pop(), Some(3));
        assert_eq!(pq.pop(), Some(4));
        assert_eq!(pq.pop(), Some(5));
        assert_eq!(pq.pop(), Some(6));
        assert_eq!(pq.pop(), None);
    }
    #[test]
    fn test_max_pq() {
        let mut pq = PQ::new_max_pq();
        pq.push(2);
        pq.push(1);
        pq.push(5);
        pq.push(3);
        pq.push(4);
        pq.push(6);
        assert_eq!(pq.peek(), Some(&6));
        assert_eq!(pq.pop(), Some(6));
        assert_eq!(pq.pop(), Some(5));
        assert_eq!(pq.pop(), Some(4));
        assert_eq!(pq.peek(), Some(&3));
        assert_eq!(pq.pop(), Some(3));
        assert_eq!(pq.pop(), Some(2));
        assert_eq!(pq.pop(), Some(1));
        assert_eq!(pq.pop(), None);
        assert_eq!(pq.pop(), None);
    }
    #[test]
    fn test_min_index_pq() {
        let mut pq = IndexPQ::new_min_pq(6);
        pq.push(0, 2);
        pq.push(1, 1);
        pq.push(2, 5);
        pq.push(3, 3);
        pq.push(4, 4);
        pq.push(5, 6);
        assert_eq!(pq.peek(), Some((1, &1)));
        assert_eq!(pq.pop(), Some((1, 1)));
        pq.change(0, 7);
        assert_eq!(pq.pop(), Some((3, 3)));
        assert_eq!(pq.pop(), Some((4, 4)));
        assert_eq!(pq.pop(), Some((2, 5)));
        assert_eq!(pq.pop(), Some((5, 6)));
        assert_eq!(pq.pop(), Some((0, 7)));
        assert_eq!(pq.pop(), None);
    }
    #[test]
    fn test_max_index_pq() {
        let mut pq = IndexPQ::new_max_pq(6);
        pq.push(0, 2);
        pq.push(1, 1);
        pq.push(2, 5);
        pq.push(3, 3);
        pq.push(4, 4);
        pq.push(5, 6);
        assert_eq!(pq.peek(), Some((5, &6)));
        assert_eq!(pq.pop(), Some((5, 6)));
        pq.change(4, 7);
        assert_eq!(pq.pop(), Some((4, 7)));
        assert_eq!(pq.pop(), Some((2, 5)));
        assert_eq!(pq.peek(), Some((3, &3)));
        assert_eq!(pq.pop(), Some((3, 3)));
        assert_eq!(pq.pop(), Some((0, 2)));
        assert_eq!(pq.pop(), Some((1, 1)));
        assert_eq!(pq.pop(), None);
        assert_eq!(pq.pop(), None);
    }
}