use std::cmp::Reverse;
use std::collections::BinaryHeap;
#[derive(Debug)]
pub struct LongHeap {
max_size: usize,
heap: BinaryHeap<Reverse<i64>>,
}
impl LongHeap {
pub fn new(max_size: usize) -> Self {
assert!(max_size >= 1, "max_size must be > 0; got: {}", max_size);
Self {
max_size,
heap: BinaryHeap::with_capacity(max_size),
}
}
pub fn new_with_initial_value(size: usize, initial_value: i64) -> Self {
let mut h = Self::new(size);
for _ in 0..size {
h.heap.push(Reverse(initial_value));
}
h
}
pub fn push(&mut self, element: i64) -> i64 {
self.heap.push(Reverse(element));
self.top()
}
pub fn insert_with_overflow(&mut self, value: i64) -> bool {
if self.heap.len() >= self.max_size {
if value < self.top() {
return false;
}
self.update_top(value);
return true;
}
self.push(value);
true
}
pub fn top(&self) -> i64 {
self.heap.peek().map_or(0, |r| r.0)
}
pub fn pop(&mut self) -> i64 {
self.heap.pop().expect("The heap is empty").0
}
pub fn update_top(&mut self, value: i64) -> i64 {
self.heap.pop();
self.heap.push(Reverse(value));
self.top()
}
pub fn size(&self) -> usize {
self.heap.len()
}
pub fn clear(&mut self) {
self.heap.clear();
}
pub fn get(&self, i: usize) -> i64 {
assert!(
i >= 1 && i <= self.heap.len(),
"index {} out of range [1, {}]",
i,
self.heap.len()
);
let slice = self.heap.iter().collect::<Vec<_>>();
slice[i - 1].0
}
}
#[cfg(test)]
mod tests {
use super::*;
use assertables::*;
#[test]
fn test_new_empty() {
let heap = LongHeap::new(10);
assert_eq!(heap.size(), 0);
}
#[test]
#[should_panic(expected = "max_size must be > 0")]
fn test_new_zero_panics() {
LongHeap::new(0);
}
#[test]
fn test_new_with_initial_value() {
let heap = LongHeap::new_with_initial_value(5, 42);
assert_eq!(heap.size(), 5);
assert_eq!(heap.top(), 42);
}
#[test]
fn test_push_and_top() {
let mut heap = LongHeap::new(10);
heap.push(5);
assert_eq!(heap.top(), 5);
heap.push(3);
assert_eq!(heap.top(), 3);
heap.push(7);
assert_eq!(heap.top(), 3);
}
#[test]
fn test_push_returns_new_top() {
let mut heap = LongHeap::new(10);
assert_eq!(heap.push(5), 5);
assert_eq!(heap.push(3), 3);
assert_eq!(heap.push(7), 3);
assert_eq!(heap.push(1), 1);
}
#[test]
fn test_pop() {
let mut heap = LongHeap::new(10);
heap.push(5);
heap.push(3);
heap.push(7);
assert_eq!(heap.pop(), 3);
assert_eq!(heap.pop(), 5);
assert_eq!(heap.pop(), 7);
assert_eq!(heap.size(), 0);
}
#[test]
#[should_panic(expected = "The heap is empty")]
fn test_pop_empty_panics() {
let mut heap = LongHeap::new(10);
heap.pop();
}
#[test]
fn test_update_top() {
let mut heap = LongHeap::new(10);
heap.push(3);
heap.push(5);
heap.push(7);
let new_top = heap.update_top(6);
assert_eq!(new_top, 5);
assert_eq!(heap.top(), 5);
}
#[test]
fn test_update_top_with_smaller() {
let mut heap = LongHeap::new(10);
heap.push(3);
heap.push(5);
heap.push(7);
let new_top = heap.update_top(1);
assert_eq!(new_top, 1);
}
#[test]
fn test_insert_with_overflow_not_full() {
let mut heap = LongHeap::new(3);
assert!(heap.insert_with_overflow(5));
assert!(heap.insert_with_overflow(3));
assert!(heap.insert_with_overflow(7));
assert_eq!(heap.size(), 3);
assert_eq!(heap.top(), 3);
}
#[test]
fn test_insert_with_overflow_full_rejected() {
let mut heap = LongHeap::new(3);
heap.push(5);
heap.push(3);
heap.push(7);
assert!(!heap.insert_with_overflow(2));
assert_eq!(heap.size(), 3);
assert_eq!(heap.top(), 3);
}
#[test]
fn test_insert_with_overflow_full_accepted() {
let mut heap = LongHeap::new(3);
heap.push(5);
heap.push(3);
heap.push(7);
assert!(heap.insert_with_overflow(6));
assert_eq!(heap.size(), 3);
assert_eq!(heap.top(), 5);
}
#[test]
fn test_size() {
let mut heap = LongHeap::new(10);
assert_eq!(heap.size(), 0);
heap.push(1);
assert_eq!(heap.size(), 1);
heap.push(2);
assert_eq!(heap.size(), 2);
heap.pop();
assert_eq!(heap.size(), 1);
}
#[test]
fn test_clear() {
let mut heap = LongHeap::new(10);
heap.push(1);
heap.push(2);
heap.push(3);
heap.clear();
assert_eq!(heap.size(), 0);
}
#[test]
fn test_get() {
let mut heap = LongHeap::new(10);
heap.push(5);
heap.push(3);
heap.push(7);
let mut elements: Vec<i64> = (1..=heap.size()).map(|i| heap.get(i)).collect();
elements.sort();
assert_eq!(elements, vec![3, 5, 7]);
}
#[test]
#[should_panic(expected = "index 0 out of range")]
fn test_get_zero_panics() {
let heap = LongHeap::new_with_initial_value(3, 1);
heap.get(0);
}
#[test]
#[should_panic(expected = "out of range")]
fn test_get_beyond_size_panics() {
let heap = LongHeap::new_with_initial_value(3, 1);
heap.get(4);
}
#[test]
fn test_min_heap_property() {
let mut heap = LongHeap::new(10);
let values = [10, 4, 15, 1, 8, 3, 12, 6];
for v in values {
heap.push(v);
}
let mut sorted = Vec::new();
while heap.size() > 0 {
sorted.push(heap.pop());
}
assert_eq!(sorted, vec![1, 3, 4, 6, 8, 10, 12, 15]);
}
#[test]
fn test_duplicate_values() {
let mut heap = LongHeap::new(10);
heap.push(5);
heap.push(5);
heap.push(5);
assert_eq!(heap.size(), 3);
assert_eq!(heap.pop(), 5);
assert_eq!(heap.pop(), 5);
assert_eq!(heap.pop(), 5);
}
#[test]
fn test_negative_values() {
let mut heap = LongHeap::new(10);
heap.push(-10);
heap.push(-5);
heap.push(-20);
assert_eq!(heap.top(), -20);
assert_eq!(heap.pop(), -20);
assert_eq!(heap.pop(), -10);
assert_eq!(heap.pop(), -5);
}
#[test]
fn test_initial_value_heap_sorted_output() {
let heap = LongHeap::new_with_initial_value(3, i64::MIN);
for i in 1..=heap.size() {
assert_eq!(heap.get(i), i64::MIN);
}
}
#[test]
fn test_grow_beyond_max_size() {
let mut heap = LongHeap::new(2);
heap.push(1);
heap.push(2);
heap.push(3);
assert_eq!(heap.size(), 3);
assert_eq!(heap.top(), 1);
}
#[test]
fn test_insert_with_overflow_boundary() {
let mut heap = LongHeap::new(3);
heap.push(3);
heap.push(5);
heap.push(7);
assert_gt!(3, heap.top() - 1);
assert!(heap.insert_with_overflow(3));
}
}