use std::fmt;
use std::iter::{FromIterator, FusedIterator};
use std::mem::{swap, ManuallyDrop};
use std::ops::{Deref, DerefMut};
use std::ptr;
pub struct WeakHeap<T> {
data: Vec<T>,
bit: Vec<bool>,
}
pub struct WeakHeapPeekMut<'a, T: 'a + Ord> {
heap: &'a mut WeakHeap<T>,
sift: bool,
}
impl<T: Ord + fmt::Debug> fmt::Debug for WeakHeapPeekMut<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("WeakHeapPeekMut")
.field(&self.heap.data[0])
.finish()
}
}
impl<T: Ord> Drop for WeakHeapPeekMut<'_, T> {
fn drop(&mut self) {
if self.sift {
unsafe { self.heap.sift_down(0) };
}
}
}
impl<T: Ord> Deref for WeakHeapPeekMut<'_, T> {
type Target = T;
fn deref(&self) -> &T {
debug_assert!(!self.heap.is_empty());
unsafe { self.heap.data.get_unchecked(0) }
}
}
impl<T: Ord> DerefMut for WeakHeapPeekMut<'_, T> {
fn deref_mut(&mut self) -> &mut T {
debug_assert!(!self.heap.is_empty());
self.sift = true;
unsafe { self.heap.data.get_unchecked_mut(0) }
}
}
impl<'a, T: Ord> WeakHeapPeekMut<'a, T> {
pub fn pop(mut this: WeakHeapPeekMut<'a, T>) -> T {
let value = this.heap.pop().unwrap();
this.sift = false;
value
}
}
impl<T: Clone> Clone for WeakHeap<T> {
fn clone(&self) -> Self {
WeakHeap {
data: self.data.clone(),
bit: self.bit.clone(),
}
}
fn clone_from(&mut self, source: &Self) {
self.data.clone_from(&source.data);
self.bit.clone_from(&source.bit);
}
}
impl<T: Ord> Default for WeakHeap<T> {
#[inline]
fn default() -> WeakHeap<T> {
WeakHeap::new()
}
}
impl<T: fmt::Debug> fmt::Debug for WeakHeap<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_list()
.entries(self.data.iter().zip(self.bit.iter()))
.finish()
}
}
impl<T: Ord> WeakHeap<T> {
#[must_use]
pub fn new() -> WeakHeap<T> {
WeakHeap {
data: vec![],
bit: vec![],
}
}
#[must_use]
pub fn with_capacity(capacity: usize) -> WeakHeap<T> {
WeakHeap {
data: Vec::with_capacity(capacity),
bit: Vec::with_capacity(capacity),
}
}
pub fn peek_mut(&mut self) -> Option<WeakHeapPeekMut<'_, T>> {
if self.is_empty() {
None
} else {
Some(WeakHeapPeekMut {
heap: self,
sift: false,
})
}
}
pub fn pop(&mut self) -> Option<T> {
self.bit.pop();
self.data.pop().map(|mut item| {
if !self.is_empty() {
swap(&mut item, &mut self.data[0]);
unsafe { self.sift_down(0) };
}
item
})
}
pub fn push(&mut self, item: T) {
let old_len = self.len();
self.data.push(item);
self.bit.push(false);
if old_len != 0 {
unsafe { self.sift_up_push(0, old_len) };
}
}
pub fn pushpop(&mut self, mut item: T) -> T {
if self.len() == 0 {
return item;
}
if self.data[0] < item {
item
} else {
swap(&mut item, &mut self.data[0]);
unsafe {
self.sift_down(0);
}
item
}
}
#[must_use = "`self` will be dropped if the result is not used"]
pub fn into_sorted_vec(mut self) -> Vec<T> {
let mut end = self.len();
while end > 1 {
end -= 1;
unsafe {
let ptr = self.data.as_mut_ptr();
std::ptr::swap(ptr, ptr.add(end));
}
unsafe { self.sift_down_range(0, end) };
}
self.into_vec()
}
unsafe fn sift_up(&mut self, start: usize, pos: usize) {
let len = self.data.len();
let mut cur = pos;
let mut ancestor = cur / 2;
while ancestor > start && (cur % 2 == *self.bit.get_unchecked(ancestor) as usize) {
cur /= 2;
ancestor /= 2;
}
if self.data.get_unchecked(ancestor) < self.data.get_unchecked(pos) {
if 2 * pos - 1 < len {
*self.bit.get_unchecked_mut(pos) ^= true;
}
let ptr = self.data.as_mut_ptr();
std::ptr::swap_nonoverlapping(ptr.add(ancestor), ptr.add(pos), 1);
}
}
unsafe fn sift_up_push(&mut self, start: usize, pos: usize) -> usize {
let len = self.data.len();
let mut hole = Hole::new(&mut self.data, pos);
let mut cur = pos;
while cur > start {
let mut ancestor = cur / 2;
while ancestor > start && (cur % 2 == *self.bit.get_unchecked(ancestor) as usize) {
cur /= 2;
ancestor /= 2;
}
if hole.get(ancestor) < hole.element() {
if 2 * pos - 1 < len {
*self.bit.get_unchecked_mut(pos) ^= true;
}
hole.move_to(ancestor);
} else {
break; }
cur = ancestor;
}
hole.pos()
}
unsafe fn sift_down_range(&mut self, start: usize, end: usize) {
if end == 1 {
return;
}
let mut pos = start.max(1);
while pos * 2 + (*self.bit.get_unchecked(pos) as usize) < end {
pos = 2 * pos + (*self.bit.get_unchecked(pos) as usize);
}
while pos > start {
if self.data.get_unchecked(start) < self.data.get_unchecked(pos) {
*self.bit.get_unchecked_mut(pos) ^= true;
let ptr = self.data.as_mut_ptr();
std::ptr::swap_nonoverlapping(ptr.add(start), ptr.add(pos), 1);
}
pos /= 2;
}
}
unsafe fn sift_down(&mut self, pos: usize) {
let len = self.len();
self.sift_down_range(pos, len);
}
fn rebuild(&mut self) {
for n in (1..self.len()).rev() {
unsafe {
self.sift_up(0, n);
}
}
}
fn rebuild_tail(&mut self, start: usize) {
if start == self.len() {
return;
}
for i in start..self.len() {
unsafe {
self.sift_up_push(0, i);
}
}
}
pub fn append(&mut self, other: &mut Self) {
if self.len() < other.len() {
swap(self, other);
}
let start = self.data.len();
self.data.append(&mut other.data);
self.bit.append(&mut other.bit);
self.rebuild_tail(start);
}
pub fn append_vec(&mut self, other: &mut Vec<T>) {
let start = self.len();
self.bit.append(&mut vec![false; other.len()]);
self.data.append(other);
self.rebuild_tail(start);
}
}
impl<T> WeakHeap<T> {
pub fn iter(&self) -> Iter<'_, T> {
Iter {
iter: self.data.iter(),
}
}
#[must_use]
pub fn peek(&self) -> Option<&T> {
self.data.get(0)
}
#[must_use]
pub fn capacity(&self) -> usize {
self.data.capacity()
}
pub fn reserve_exact(&mut self, additional: usize) {
self.data.reserve_exact(additional);
self.bit.reserve_exact(additional);
}
pub fn reserve(&mut self, additional: usize) {
self.data.reserve(additional);
self.bit.reserve(additional);
}
pub fn shrink_to_fit(&mut self) {
self.data.shrink_to_fit();
self.bit.shrink_to_fit();
}
#[inline]
pub fn shrink_to(&mut self, min_capacity: usize) {
self.data.shrink_to(min_capacity);
self.bit.shrink_to(min_capacity);
}
#[must_use = "`self` will be dropped if the result is not used"]
pub fn into_vec(self) -> Vec<T> {
self.data
}
#[must_use]
pub fn len(&self) -> usize {
self.data.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
#[inline]
pub fn drain(&mut self) -> Drain<'_, T> {
self.bit.clear();
Drain {
iter: self.data.drain(..),
}
}
pub fn clear(&mut self) {
self.drain();
}
}
struct Hole<'a, T: 'a> {
data: &'a mut [T],
elt: ManuallyDrop<T>,
pos: usize,
}
impl<'a, T> Hole<'a, T> {
#[inline]
unsafe fn new(data: &'a mut [T], pos: usize) -> Self {
debug_assert!(pos < data.len());
let elt = ptr::read(data.get_unchecked(pos));
Hole {
data,
elt: ManuallyDrop::new(elt),
pos,
}
}
#[inline]
fn pos(&self) -> usize {
self.pos
}
#[inline]
fn element(&self) -> &T {
&self.elt
}
#[inline]
unsafe fn get(&self, index: usize) -> &T {
debug_assert!(index != self.pos);
debug_assert!(index < self.data.len());
self.data.get_unchecked(index)
}
#[inline]
unsafe fn move_to(&mut self, index: usize) {
debug_assert!(index != self.pos);
debug_assert!(index < self.data.len());
let ptr = self.data.as_mut_ptr();
let index_ptr: *const _ = ptr.add(index);
let hole_ptr = ptr.add(self.pos);
ptr::copy_nonoverlapping(index_ptr, hole_ptr, 1);
self.pos = index;
}
}
impl<T> Drop for Hole<'_, T> {
#[inline]
fn drop(&mut self) {
unsafe {
let pos = self.pos;
ptr::copy_nonoverlapping(&*self.elt, self.data.get_unchecked_mut(pos), 1);
}
}
}
impl<T: Ord> From<Vec<T>> for WeakHeap<T> {
fn from(vec: Vec<T>) -> WeakHeap<T> {
let n = vec.len();
let mut heap = WeakHeap {
data: vec,
bit: vec![false; n],
};
heap.rebuild();
heap
}
}
impl<T: Ord, const N: usize> From<[T; N]> for WeakHeap<T> {
fn from(arr: [T; N]) -> Self {
arr.into_iter().collect()
}
}
impl<T> From<WeakHeap<T>> for Vec<T> {
fn from(heap: WeakHeap<T>) -> Vec<T> {
heap.data
}
}
impl<T: Ord> FromIterator<T> for WeakHeap<T> {
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> WeakHeap<T> {
WeakHeap::from(iter.into_iter().collect::<Vec<_>>())
}
}
impl<T: Ord> Extend<T> for WeakHeap<T> {
fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
for x in iter {
self.push(x);
}
}
}
impl<'a, T: 'a + Ord + Copy> Extend<&'a T> for WeakHeap<T> {
fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
self.extend(iter.into_iter().cloned());
}
}
impl<T> IntoIterator for WeakHeap<T> {
type Item = T;
type IntoIter = IntoIter<T>;
fn into_iter(self) -> IntoIter<T> {
IntoIter {
iter: self.data.into_iter(),
}
}
}
impl<'a, T> IntoIterator for &'a WeakHeap<T> {
type Item = &'a T;
type IntoIter = Iter<'a, T>;
fn into_iter(self) -> Iter<'a, T> {
self.iter()
}
}
#[derive(Clone)]
pub struct Iter<'a, T: 'a> {
iter: std::slice::Iter<'a, T>,
}
impl<T: fmt::Debug> fmt::Debug for Iter<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("Iter").field(&self.iter.as_slice()).finish()
}
}
impl<'a, T> Iterator for Iter<'a, T> {
type Item = &'a T;
#[inline]
fn next(&mut self) -> Option<&'a T> {
self.iter.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[inline]
fn last(self) -> Option<&'a T> {
self.iter.last()
}
}
impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
#[inline]
fn next_back(&mut self) -> Option<&'a T> {
self.iter.next_back()
}
}
impl<T> FusedIterator for Iter<'_, T> {}
#[derive(Clone)]
pub struct IntoIter<T> {
iter: std::vec::IntoIter<T>,
}
impl<T: fmt::Debug> fmt::Debug for IntoIter<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("IntoIter")
.field(&self.iter.as_slice())
.finish()
}
}
impl<T> Iterator for IntoIter<T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
self.iter.next()
}
}
impl<T> DoubleEndedIterator for IntoIter<T> {
#[inline]
fn next_back(&mut self) -> Option<T> {
self.iter.next_back()
}
}
impl<T> FusedIterator for IntoIter<T> {}
#[derive(Debug)]
pub struct Drain<'a, T: 'a> {
iter: std::vec::Drain<'a, T>,
}
impl<T> Iterator for Drain<'_, T> {
type Item = T;
#[inline]
fn next(&mut self) -> Option<T> {
self.iter.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
}
impl<T> DoubleEndedIterator for Drain<'_, T> {
#[inline]
fn next_back(&mut self) -> Option<T> {
self.iter.next_back()
}
}
impl<T> FusedIterator for Drain<'_, T> {}
#[cfg(test)]
mod tests;