cp-library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub kobejean/cp-library

:heavy_check_mark: cp_library/ds/heap/min_k_heap_cls.py

Depends on

Verified with

Code

import cp_library.ds.heap.__header__
from typing import Iterable
from cp_library.ds.heap.max_heap_cls import MaxHeap
from cp_library.ds.heap.k_heap_mixin import KHeapMixin
from cp_library.misc.typing import _T

class MinKHeap(KHeapMixin[_T], MaxHeap[_T]):
    '''MinKHeap[K: int, T: type, N: Union[int,None]]'''
    def __init__(self, K: int, iterable: Iterable[_T] = None):
        MaxHeap.__init__(self, iterable)
        KHeapMixin.__init__(self, K)
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from typing import Iterable
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 MaxHeap(HeapProtocol[_T], UserList[_T]):
    def __init__(self, iterable: Iterable[_T] = None):
        super().__init__(iterable)
        heapify_max(self.data)
    def pop(self): return heappop_max(self.data)
    def push(self, item: _T): heappush_max(self.data, item)
    def pushpop(self, item: _T): return heappushpop_max(self.data, item)
    def replace(self, item: _T): return heapreplace_max(self.data, item)
from typing import Union


import typing
from collections import deque
from numbers import Number
from types import GenericAlias 
from typing import Callable, Collection, Iterator, Union
import os
import sys
from io import BytesIO, IOBase


class FastIO(IOBase):
    BUFSIZE = 8192
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = BytesIO()
        self.writable = "x" in file.mode or "r" not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        BUFSIZE = self.BUFSIZE
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        BUFSIZE = self.BUFSIZE
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b"\n") + (not b)
            ptr = self.buffer.tell()
            self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr)
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            self.buffer.truncate(0), self.buffer.seek(0)


class IOWrapper(IOBase):
    stdin: 'IOWrapper' = None
    stdout: 'IOWrapper' = None
    
    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable

    def write(self, s):
        return self.buffer.write(s.encode("ascii"))
    
    def read(self):
        return self.buffer.read().decode("ascii")
    
    def readline(self):
        return self.buffer.readline().decode("ascii")

sys.stdin = IOWrapper.stdin = IOWrapper(sys.stdin)
sys.stdout = IOWrapper.stdout = IOWrapper(sys.stdout)

class TokenStream(Iterator):
    stream = IOWrapper.stdin

    def __init__(self):
        self.queue = deque()

    def __next__(self):
        if not self.queue: self.queue.extend(self._line())
        return self.queue.popleft()
    
    def wait(self):
        if not self.queue: self.queue.extend(self._line())
        while self.queue: yield
 
    def _line(self):
        return TokenStream.stream.readline().split()

    def line(self):
        if self.queue:
            A = list(self.queue)
            self.queue.clear()
            return A
        return self._line()
TokenStream.default = TokenStream()

class CharStream(TokenStream):
    def _line(self):
        return TokenStream.stream.readline().rstrip()
CharStream.default = CharStream()


ParseFn = Callable[[TokenStream],_T]
class Parser:
    def __init__(self, spec: Union[type[_T],_T]):
        self.parse = Parser.compile(spec)

    def __call__(self, ts: TokenStream) -> _T:
        return self.parse(ts)
    
    @staticmethod
    def compile_type(cls: type[_T], args = ()) -> _T:
        if issubclass(cls, Parsable):
            return cls.compile(*args)
        elif issubclass(cls, (Number, str)):
            def parse(ts: TokenStream): return cls(next(ts))              
            return parse
        elif issubclass(cls, tuple):
            return Parser.compile_tuple(cls, args)
        elif issubclass(cls, Collection):
            return Parser.compile_collection(cls, args)
        elif callable(cls):
            def parse(ts: TokenStream):
                return cls(next(ts))              
            return parse
        else:
            raise NotImplementedError()
    
    @staticmethod
    def compile(spec: Union[type[_T],_T]=int) -> ParseFn[_T]:
        if isinstance(spec, (type, GenericAlias)):
            cls = typing.get_origin(spec) or spec
            args = typing.get_args(spec) or tuple()
            return Parser.compile_type(cls, args)
        elif isinstance(offset := spec, Number): 
            cls = type(spec)  
            def parse(ts: TokenStream): return cls(next(ts)) + offset
            return parse
        elif isinstance(args := spec, tuple):      
            return Parser.compile_tuple(type(spec), args)
        elif isinstance(args := spec, Collection):  
            return Parser.compile_collection(type(spec), args)
        elif isinstance(fn := spec, Callable): 
            def parse(ts: TokenStream): return fn(next(ts))
            return parse
        else:
            raise NotImplementedError()

    @staticmethod
    def compile_line(cls: _T, spec=int) -> ParseFn[_T]:
        if spec is int:
            fn = Parser.compile(spec)
            def parse(ts: TokenStream): return cls([int(token) for token in ts.line()])
            return parse
        else:
            fn = Parser.compile(spec)
            def parse(ts: TokenStream): return cls([fn(ts) for _ in ts.wait()])
            return parse

    @staticmethod
    def compile_repeat(cls: _T, spec, N) -> ParseFn[_T]:
        fn = Parser.compile(spec)
        def parse(ts: TokenStream): return cls([fn(ts) for _ in range(N)])
        return parse

    @staticmethod
    def compile_children(cls: _T, specs) -> ParseFn[_T]:
        fns = tuple((Parser.compile(spec) for spec in specs))
        def parse(ts: TokenStream): return cls([fn(ts) for fn in fns])  
        return parse
            
    @staticmethod
    def compile_tuple(cls: type[_T], specs) -> ParseFn[_T]:
        if isinstance(specs, (tuple,list)) and len(specs) == 2 and specs[1] is ...:
            return Parser.compile_line(cls, specs[0])
        else:
            return Parser.compile_children(cls, specs)

    @staticmethod
    def compile_collection(cls, specs):
        if not specs or len(specs) == 1 or isinstance(specs, set):
            return Parser.compile_line(cls, *specs)
        elif (isinstance(specs, (tuple,list)) and len(specs) == 2 and isinstance(specs[1], int)):
            return Parser.compile_repeat(cls, specs[0], specs[1])
        else:
            raise NotImplementedError()

class Parsable:
    @classmethod
    def compile(cls):
        def parser(ts: TokenStream): return cls(next(ts))
        return parser

class KHeapMixin(HeapProtocol[_T], Parsable):
    '''KHeapMixin[K: int, T: type, N: Union[int,None]]'''
    def __init__(heap, K: int):
        heap.K = K

    def added(heap, item: _T): ...

    def removed(heap, item: _T): ...
    
    def pop(heap):
        item = super().pop()
        heap.removed(item)
        return item
    
    def push(heap, item: _T):
        if len(heap) < heap._K:
            heap.added(item)
            super().push(item)
        elif heap._K:
            assert len(heap) == heap._K, f'{len(heap)=} {heap._K}'
            heap.pushpop(item)
    
    def pushpop(heap, item: _T):
        if item != (remove := super().pushpop(item)):
            heap.removed(remove)
            heap.added(item)
            return remove
        else:
            return item
    
    def replace(heap, item: _T):
        remove = super().replace(item)
        heap.removed(remove)
        heap.added(item)
        return remove
    
    
    @property
    def K(heap):
        return heap._K

    @K.setter
    def K(heap, K):
        heap._K = K
        if K is not None:
            while len(heap) > K:
                heap.pop()
    
    @classmethod
    def compile(cls, K: int, T: type, N: Union[int,None] = None):
        elm = Parser.compile(T)
        if N is None:
            def parse(ts: TokenStream):
                return cls(K, (elm(ts) for _ in ts.wait()))
        else:
            def parse(ts: TokenStream):
                return cls(K, (elm(ts) for _ in range(N)))
        return parse

class MinKHeap(KHeapMixin[_T], MaxHeap[_T]):
    '''MinKHeap[K: int, T: type, N: Union[int,None]]'''
    def __init__(self, K: int, iterable: Iterable[_T] = None):
        MaxHeap.__init__(self, iterable)
        KHeapMixin.__init__(self, K)
Back to top page