cp_library/ds/heap/min_k_heap_cls.py

• View this file on GitHub
• Last update: 2025-08-24 20:27:05+09:00

Depends on

• cp_library/ds/heap/heap_base_cls.py
• cp_library/ds/heap/heapify_max_fn.py
• cp_library/ds/heap/heappop_max_fn.py
• cp_library/ds/heap/heappush_max_fn.py
• cp_library/ds/heap/heappushpop_max_fn.py
• cp_library/ds/heap/heapreplace_max_fn.py
• cp_library/ds/heap/heapsiftdown_max_fn.py
• cp_library/ds/heap/heapsiftup_max_fn.py
• cp_library/ds/heap/k_heap_mixin.py
• cp_library/ds/heap/max_heap_cls.py
• cp_library/io/io_base_cls.py
• cp_library/io/parsable_cls.py
• cp_library/io/parser_cls.py

Verified with

• test/atcoder/abc/abc249_f_min_k_heap.test.py

Code

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

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 typing import TypeVar

_S = TypeVar('S'); _T = TypeVar('T'); _U = TypeVar('U'); _T1 = TypeVar('T1'); _T2 = TypeVar('T2'); _T3 = TypeVar('T3'); _T4 = TypeVar('T4'); _T5 = TypeVar('T5'); _T6 = TypeVar('T6')



def heapsiftup_max(heap: list, pos: int):
    n, item, c = len(heap)-1, heap[pos], pos<<1|1
    while c < n and item < heap[c := c+(heap[c]<heap[c+1])]: heap[pos], pos, c = heap[c], c, c<<1|1
    if c == n and item < heap[c]: heap[pos], pos = heap[c], c
    heap[pos] = item

def heapify_max(x: list):
    for i in reversed(range(len(x)//2)): heapsiftup_max(x, i)

def heappop_max(heap: list):
    item = heap.pop()
    if heap: item, heap[0] = heap[0], item; heapsiftup_max(heap, 0)
    return item

def heapsiftdown_max(heap: list, root: int, pos: int):
    item = heap[pos]
    while root < pos and heap[p := (pos-1)>>1] < item: heap[pos], pos = heap[p], p
    heap[pos] = item

def heappush_max(heap: list, item):
    heap.append(item); heapsiftdown_max(heap, 0, len(heap)-1)

def heappushpop_max(heap: list, item):
    if heap and heap[0] > item: item, heap[0] = heap[0], item; heapsiftup_max(heap, 0)
    return item

def heapreplace_max(heap: list, item):
    item, heap[0] = heap[0], item; heapsiftup_max(heap, 0)
    return item
from typing import Generic

class HeapBase(Generic[_T]):
    def peek(heap) -> _T: return heap.data[0]
    def pop(heap) -> _T: ...
    def push(heap, item: _T): ...
    def pushpop(heap, item: _T) -> _T: ...
    def replace(heap, item: _T) -> _T: ...
    def __contains__(heap, item: _T): return item in heap.data
    def __len__(heap): return len(heap.data)
    def clear(heap): heap.data.clear()

class MaxHeap(HeapBase[_T]):
    def __init__(self, iterable: Iterable[_T] = None): self.data = list(iterable) if iterable else []; 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
from types import GenericAlias


class Parsable:
    @classmethod
    def compile(cls):
        def parser(io: 'IOBase'): return cls(next(io))
        return parser
    @classmethod
    def __class_getitem__(cls, item): return GenericAlias(cls, item)

class KHeapMixin(HeapBase[_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: 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(io: IOBase): return cls(K, (elm(io) for _ in io.wait()))
        else:
            def parse(io: IOBase): return cls(K, (elm(io) for _ in range(N)))
        return parse

class IOBase:
    @property
    def char(io) -> bool: ...
    @property
    def writable(io) -> bool: ...
    def __next__(io) -> str: ...
    def write(io, s: str) -> None: ...
    def readline(io) -> str: ...
    def readtoken(io) -> str: ...
    def readtokens(io) -> list[str]: ...
    def readints(io) -> list[int]: ...
    def readdigits(io) -> list[int]: ...
    def readnums(io) -> list[int]: ...
    def readchar(io) -> str: ...
    def readchars(io) -> str: ...
    def readinto(io, lst: list[str]) -> list[str]: ...
    def readcharsinto(io, lst: list[str]) -> list[str]: ...
    def readtokensinto(io, lst: list[str]) -> list[str]: ...
    def readintsinto(io, lst: list[int]) -> list[int]: ...
    def readdigitsinto(io, lst: list[int]) -> list[int]: ...
    def readnumsinto(io, lst: list[int]) -> list[int]: ...
    def wait(io): ...
    def flush(io) -> None: ...
    def line(io) -> list[str]: ...
import typing
from numbers import Number
from typing import Callable, Collection

class Parser:
    def __init__(self, spec):  self.parse = Parser.compile(spec)
    def __call__(self, io: IOBase): return self.parse(io)
    @staticmethod
    def compile_type(cls, args = ()):
        if issubclass(cls, Parsable): return cls.compile(*args)
        elif issubclass(cls, (Number, str)):
            def parse(io: IOBase): return cls(next(io))              
            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(io: IOBase): return cls(next(io))              
            return parse
        else: raise NotImplementedError()
    @staticmethod
    def compile(spec=int):
        if isinstance(spec, (type, GenericAlias)):
            cls, args = typing.get_origin(spec) or spec, typing.get_args(spec) or tuple()
            return Parser.compile_type(cls, args)
        elif isinstance(offset := spec, Number): 
            cls = type(spec)  
            def parse(io: IOBase): return cls(next(io)) + 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(io: IOBase): return fn(next(io))
            return parse
        else: raise NotImplementedError()
    @staticmethod
    def compile_line(cls, spec=int):
        if spec is int:
            def parse(io: IOBase): return cls(io.readnums())
        elif spec is str:
            def parse(io: IOBase): return cls(io.line())
        else:
            fn = Parser.compile(spec)
            def parse(io: IOBase): return cls((fn(io) for _ in io.wait()))
        return parse
    @staticmethod
    def compile_repeat(cls, spec, N):
        fn = Parser.compile(spec)
        def parse(io: IOBase): return cls([fn(io) for _ in range(N)])
        return parse
    @staticmethod
    def compile_children(cls, specs):
        fns = tuple((Parser.compile(spec) for spec in specs))
        def parse(io: IOBase): return cls([fn(io) for fn in fns])  
        return parse
    @staticmethod
    def compile_tuple(cls, specs):
        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 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