cp_library/ds/heap/max_priority_queue_cls.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00

Depends on

• cp_library/ds/heap/heap_proto.py
• cp_library/ds/heap/heapq_max_import.py

Code

import cp_library.ds.heap.__header__

from collections import UserList
from cp_library.ds.heap.heapq_max_import import heapify_max, heappop_max, heappush_max, heappushpop_max, heapreplace_max
from cp_library.ds.heap.heap_proto import HeapProtocol

class MaxPriorityQueue(HeapProtocol[int], UserList[int]):
    
    def __init__(self, N: int, ids: list[int] = None, priorities: list[int] = None, /):
        self.shift = N.bit_length()
        self.mask = (1 << self.shift)-1
        if ids is None:
            super().__init__()
        elif priorities is None:
            heapify_max(ids)
            self.data = ids
        else:
            M = len(ids)
            data = [0]*M
            for i in range(M):
                data[i] = self.encode(ids[i], priorities[i]) 
            heapify_max(data)
            self.data = data

    def encode(self, id, priority):
        return priority << self.shift | id
    
    def decode(self, encoded):
        return self.mask & encoded, encoded >> self.shift
    
    def pop(self):
        return self.decode(heappop_max(self.data))
    
    def push(self, id: int, priority: int):
        heappush_max(self.data, self.encode(id, priority))

    def pushpop(self, id: int, priority: int):
        return self.decode(heappushpop_max(self.data, self.encode(id, priority)))
    
    def replace(self, id: int, priority: int):
        return self.decode(heapreplace_max(self.data, self.encode(id, priority)))

    def peek(self):
        return self.decode(self.data[0])

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''

from collections import UserList
from typing import TypeVar
_T = TypeVar('T')
def heappop_max(heap: list[_T], /) -> _T: ...
def heapsiftdown_max(heap: list[_T], root: int, pos: int): ...
def heapsiftup_max(heap: list[_T], pos: int): ...
def heapsiftdown(heap: list[_T], root: int, pos: int): ...
def heapsiftup(heap: list[_T], pos: int): ...

from heapq import (
    _heapify_max as heapify_max, 
    _heappop_max as heappop_max, 
    _siftdown_max as heapsiftdown_max,
    _siftup_max as heapsiftup_max,
    _siftdown as heapsiftdown,
    _siftup as heapsiftup
)

def heappush_max(heap: list[_T], item: _T):
    '''Push item onto heap, maintaining the heap invariant.'''
    heap.append(item)
    heapsiftdown_max(heap, 0, len(heap)-1)

def heapreplace_max(heap: list[_T], item: _T) -> _T:
    '''Pop and return the current largest value, and add the new item.

    This is more efficient than heappop_max() followed by heappush_max(), and can be
    more appropriate when using a fixed-size heap.  Note that the value
    returned may be larger than item!  That constrains reasonable uses of
    this routine unless written as part of a conditional replacement:

        if item > heap[0]:
            item = heapreplace_max(heap, item)
    '''
    returnitem = heap[0]
    heap[0] = item
    heapsiftup_max(heap, 0)
    return returnitem

def heappushpop_max(heap: list[_T], item: _T) -> _T:
    '''Fast version of a heappush_max followed by a heappop_max.'''
    if heap and heap[0] > item:
        item, heap[0] = heap[0], item
        heapsiftup_max(heap, 0)
    return item

from typing import Generic

class HeapProtocol(Generic[_T]):
    def pop(self) -> _T: ...
    def push(self, item: _T): ...
    def pushpop(self, item: _T) -> _T: ...
    def replace(self, item: _T) -> _T: ...

class MaxPriorityQueue(HeapProtocol[int], UserList[int]):
    
    def __init__(self, N: int, ids: list[int] = None, priorities: list[int] = None, /):
        self.shift = N.bit_length()
        self.mask = (1 << self.shift)-1
        if ids is None:
            super().__init__()
        elif priorities is None:
            heapify_max(ids)
            self.data = ids
        else:
            M = len(ids)
            data = [0]*M
            for i in range(M):
                data[i] = self.encode(ids[i], priorities[i]) 
            heapify_max(data)
            self.data = data

    def encode(self, id, priority):
        return priority << self.shift | id
    
    def decode(self, encoded):
        return self.mask & encoded, encoded >> self.shift
    
    def pop(self):
        return self.decode(heappop_max(self.data))
    
    def push(self, id: int, priority: int):
        heappush_max(self.data, self.encode(id, priority))

    def pushpop(self, id: int, priority: int):
        return self.decode(heappushpop_max(self.data, self.encode(id, priority)))
    
    def replace(self, id: int, priority: int):
        return self.decode(heapreplace_max(self.data, self.encode(id, priority)))

    def peek(self):
        return self.decode(self.data[0])

Back to top page