可删除元素的优先队列

WARNING

需要用到 usePriorityQueue,请到 优先队列 PriorityQueue 复制代码

TypeScript
type Comparator<T> = (o1: T, o2: T) => number

export function useRemovablePriorityQueue<T>(
    compare: Comparator<T> = (a, b) => (a === b ? 0 : a > b ? 1 : -1),
) {
    const pq = usePriorityQueue(compare)
    const map = new Map<T, number>()
    let _size = 0

    const removeUnusable = () => {
        while (!pq.isEmpty() && !map.has(pq.peek())) pq.pop()
    }

    const del = (e: T) => {
        if (!map.has(e)) return false

        map.set(e, map.get(e) - 1)
        if (map.get(e) <= 0) map.delete(e)

        return true
    }

    const push = (e: T) => {
        pq.push(e)

        map.set(e, (map.get(e) || 0) + 1)
        _size++
    }

    const peek = () => {
        removeUnusable()
        return pq.peek()
    }

    const pop = () => {
        removeUnusable()
        if (pq.isEmpty()) return null

        const res = pq.pop()
        _size--
        del(res)
        return res
    }

    const remove = (e: T) => {
        const hasElement = del(e)
        if (hasElement) _size--

        return hasElement
    }

    const size = () => _size

    const isEmpty = () => size() === 0

    return { push, pop, isEmpty, size, remove, peek }
}
JavaScript
export function useRemovablePriorityQueue(
    compare = (a, b) => (a === b ? 0 : a > b ? 1 : -1),
) {
    const pq = usePriorityQueue(compare)
    const map = new Map()
    let _size = 0
    const removeUnusable = () => {
        while (!pq.isEmpty() && !map.has(pq.peek())) pq.pop()
    }
    const del = (e) => {
        if (!map.has(e)) return false
        map.set(e, map.get(e) - 1)
        if (map.get(e) <= 0) map.delete(e)
        return true
    }
    const push = (e) => {
        pq.push(e)
        map.set(e, (map.get(e) || 0) + 1)
        _size++
    }
    const peek = () => {
        removeUnusable()
        return pq.peek()
    }
    const pop = () => {
        removeUnusable()
        if (pq.isEmpty()) return null
        const res = pq.pop()
        _size--
        del(res)
        return res
    }
    const remove = (e) => {
        const hasElement = del(e)
        if (hasElement) _size--
        return hasElement
    }
    const size = () => _size
    const isEmpty = () => size() === 0
    return { push, pop, isEmpty, size, remove, peek }
}
Rust
use crate::priority_queue::PriorityQueue;

pub struct RemovablePriorityQueue<T>
where
    T: Eq + std::hash::Hash + Copy + Ord,
{
    pq: PriorityQueue<T>,
    exist: std::collections::HashMap<T, i32>,
    size: usize,
}

impl<T> RemovablePriorityQueue<T>
where
    T: Eq + std::hash::Hash + Copy + Ord,
{
    pub fn new() -> Self {
        let pq = PriorityQueue::new();
        let exist = std::collections::HashMap::new();
        Self { pq, exist, size: 0 }
    }

    pub fn push(&mut self, value: T) {
        self.pq.push(value);
        self.exist.entry(value).and_modify(|x| *x += 1).or_insert(1);
        self.size += 1;
    }

    pub fn pop(&mut self) -> Option<T> {
        self.remove_unusable();
        if self.pq.is_empty() {
            return None;
        }
        let result = self.pq.pop().unwrap();
        if let Some(x) = self.exist.get_mut(&result) {
            *x -= 1;
            if *x == 0 {
                self.exist.remove(&result);
            }
            self.size -= 1;
        }
        Some(result)
    }

    pub fn peek(&mut self) -> Option<&T> {
        self.remove_unusable();
        self.pq.peek()
    }

    pub fn remove(&mut self, value: T) {
        if let Some(x) = self.exist.get_mut(&value) {
            *x -= 1;
            if *x == 0 {
                self.exist.remove(&value);
            }
            self.size -= 1;
        }
    }

    fn remove_unusable(&mut self) {
        while !self.pq.is_empty() && !self.exist.contains_key(self.pq.peek().unwrap()) {
            self.pq.pop();
        }
    }

    pub fn len(&self) -> usize {
        self.size
    }

    pub fn is_empty(&self) -> bool {
        self.size == 0
    }

    pub fn clear(&mut self) {
        self.pq.clear();
        self.exist.clear();
        self.size = 0;
    }
}
Go
package algo

type RemovablePriorityQueue[T comparable] struct {
    pq    *PriorityQueue[T]
    exist map[T]int
    size  int
}

func NewRemovablePriorityQueue[T comparable](compare func(a, b T) int) *RemovablePriorityQueue[T] {
    return &RemovablePriorityQueue[T]{
        pq:    NewPriorityQueue(compare),
        exist: make(map[T]int),
        size:  0,
    }
}

func (rpq *RemovablePriorityQueue[T]) removeUnusable() {
    for !rpq.pq.IsEmpty() {
        if peek, ok := rpq.pq.Peek(); ok {
            if _, exists := rpq.exist[peek]; exists {
                break
            }
            rpq.pq.Pop()
        } else {
            break
        }
    }
}

func (rpq *RemovablePriorityQueue[T]) Push(v T) {
    rpq.pq.Push(v)
    rpq.exist[v]++
    rpq.size++
}

func (rpq *RemovablePriorityQueue[T]) Pop() (T, bool) {
    rpq.removeUnusable()

    if rpq.pq.IsEmpty() {
        var zero T
        return zero, false
    }

    res, _ := rpq.pq.Pop()
    rpq.size--
    rpq.del(res)
    return res, true
}

func (rpq *RemovablePriorityQueue[T]) Peek() (T, bool) {
    rpq.removeUnusable()
    return rpq.pq.Peek()
}

func (rpq *RemovablePriorityQueue[T]) Remove(v T) bool {
    if _, exists := rpq.exist[v]; !exists {
        return false
    }

    rpq.del(v)
    rpq.size--
    return true
}

func (rpq *RemovablePriorityQueue[T]) del(v T) {
    if count, exists := rpq.exist[v]; exists {
        if count > 1 {
            rpq.exist[v] = count - 1
        } else {
            delete(rpq.exist, v)
        }
    }
}

func (rpq *RemovablePriorityQueue[T]) Len() int {
    return rpq.size
}

func (rpq *RemovablePriorityQueue[T]) IsEmpty() bool {
    return rpq.size == 0
}