use crate::collections::ItemPriQueue;
use num_traits::{FromPrimitive, ToPrimitive};
struct BinHeapItem<K, V, ID> {
key: K,
value: V,
pos: ID,
}
pub struct BinHeap<K, V, ID = u32> {
heap: Vec<ID>,
data: Vec<BinHeapItem<K, V, ID>>,
free: Option<ID>,
}
impl<K, V, ID> BinHeap<K, V, ID> {
pub fn new() -> Self {
Default::default()
}
}
impl<K, V, ID> Default for BinHeap<K, V, ID> {
fn default() -> Self {
BinHeap {
heap: vec![],
data: vec![],
free: None,
}
}
}
impl<K, V, ID> ItemPriQueue<K, V> for BinHeap<K, V, ID>
where
K: Clone,
V: PartialOrd + Clone,
ID: FromPrimitive + ToPrimitive + Copy + Eq,
{
type Item = ID;
fn clear(&mut self) {
self.heap.clear();
self.data.clear();
self.free = None;
}
fn is_empty(&self) -> bool {
self.heap.is_empty()
}
fn value(&self, item: &ID) -> &V {
&self.data[item.to_usize().unwrap()].value
}
fn push(&mut self, key: K, value: V) -> ID {
let item = if let Some(item) = self.free {
let idx = item.to_usize().unwrap();
let next = self.data[idx].pos;
if next == item {
self.free = None
} else {
self.free = Some(next)
}
self.data[idx] = BinHeapItem {
key,
value,
pos: ID::from_usize(self.heap.len()).unwrap(),
};
item
} else {
let item = ID::from_usize(self.data.len()).unwrap();
self.data.push(BinHeapItem {
key,
value,
pos: ID::from_usize(self.heap.len()).unwrap(),
});
item
};
self.heap.push(item);
self.upheap(item);
item
}
fn decrease_key(&mut self, item: &mut ID, value: V) -> bool {
let idx = item.to_usize().unwrap();
if self.data[idx].value > value {
self.data[idx].value = value;
self.upheap(*item);
true
} else {
false
}
}
fn pop_min(&mut self) -> Option<(K, V)> {
if self.heap.is_empty() {
return None;
}
let min_item = self.heap.swap_remove(0);
let min_idx = min_item.to_usize().unwrap();
if let Some(next) = self.free {
self.data[min_idx].pos = next;
} else {
self.data[min_idx].pos = min_item;
}
self.free = Some(min_item);
if !self.heap.is_empty() {
let n = self.heap.len();
let item = *self.heap.first().unwrap();
let idx = item.to_usize().unwrap();
let value = self.data[idx].value.clone();
let mut cur_pos = 0;
loop {
let left_pos = 2 * cur_pos + 1;
let right_pos = left_pos + 1;
let (next_pos, next_idx) = if left_pos >= n {
break;
} else if right_pos >= n {
(left_pos, self.heap[left_pos].to_usize().unwrap())
} else {
let left_idx = self.heap[left_pos].to_usize().unwrap();
let right_idx = self.heap[right_pos].to_usize().unwrap();
if self.data[left_idx].value < self.data[right_idx].value {
(left_pos, left_idx)
} else {
(right_pos, right_idx)
}
};
if value <= self.data[next_idx].value {
break;
}
self.heap[cur_pos] = self.heap[next_pos];
self.data[next_idx].pos = ID::from_usize(cur_pos).unwrap();
cur_pos = next_pos;
}
self.heap[cur_pos] = item;
self.data[idx].pos = ID::from_usize(cur_pos).unwrap();
}
Some((self.data[min_idx].key.clone(), self.data[min_idx].value.clone()))
}
}
impl<K, V, ID> BinHeap<K, V, ID>
where
V: PartialOrd + Clone,
ID: FromPrimitive + ToPrimitive + Clone + Eq,
{
fn upheap(&mut self, item: ID) {
let idx = item.to_usize().unwrap();
let value = self.data[idx].value.clone();
let mut cur_pos = self.data[idx].pos.to_usize().unwrap();
while cur_pos > 0 {
let parent_pos = (cur_pos - 1) / 2;
let parent_idx = self.heap[parent_pos].to_usize().unwrap();
if value > self.data[parent_idx].value {
break;
}
self.heap[cur_pos] = self.heap[parent_pos].clone();
self.data[parent_idx].pos = ID::from_usize(cur_pos).unwrap();
cur_pos = parent_pos;
}
self.data[idx].pos = ID::from_usize(cur_pos).unwrap();
self.heap[cur_pos] = item;
}
}