Heaps & Priority Queues

#Topic: Min-Heap Basics -- heappush, heappop, heapify, peek
# Time  : push O(log n), pop O(log n), heapify O(n)
#Space: O(n) -- Heap size

import heapq


def demo_basic_operations() -> None:
    """heapq basic operations demo."""
    print("=" * 50)
    print("Min-Heap Basic Operations")
    print("=" * 50)

    h: list[int] = []

    # heappush -- O(log n)
    for val in [5, 3, 8, 1, 9, 2, 7]:
        heapq.heappush(h, val)
        print(f"push({val}) -> Heap root (minimum): {h[0]}, Heap: {h}")

    print()
    #heappop -- O(log n): Always returns the minimum value
    print("pop order (always minimum value):")
    snapshot = list(h)
    while snapshot:
        val = heapq.heappop(snapshot)
        print(f"  pop() -> {val}")


def demo_heapify() -> None:
    """heapify -- Convert a list into a heap in O(n)."""
    print()
    print("=" * 50)
    print("heapify -- O(n) in-place conversion")
    print("=" * 50)

    lst = [9, 5, 3, 7, 1, 8, 2, 6, 4]
    print("Before conversion:", lst)

    heapq.heapify(lst)   #In-place, no return value
    print("After conversion:", lst)
    print("Root (minimum):", lst[0])

    #Note: heapify returns None
    lst2 = [3, 1, 2]
    ret = heapq.heapify(lst2)
    print(f"\nheapify() return value: {ret}  (None -- no return value)")
    print("After converting lst2:", lst2)


def demo_peek() -> None:
    """peek -- View the heap root (without removing)."""
    print()
    print("=" * 50)
    print("peek -- O(1) check minimum value")
    print("=" * 50)

    h = [4, 7, 2, 9, 1, 5]
    heapq.heapify(h)
    print("Heap:", h)
    print("peek (h[0]):", h[0])   #Check the minimum value without removal
    print("No heap change:", h)


def demo_tuple_heap() -> None:
    """Tuple heap -- (priority, value) pattern."""
    print()
    print("=" * 50)
    print("Tuple heap -- Priority queue")
    print("=" * 50)

    #Task scheduling: (priority, task name)
    h: list[tuple[int, str]] = []
    tasks = [(3, "Low priority"), (1, "Urgent"), (2, "Normal"), (1, "Urgent2")]
    for priority, name in tasks:
        heapq.heappush(h, (priority, name))
        print(f"  push({priority}, {name!r})")

    print("\nProcessing order:")
    while h:
        priority, name = heapq.heappop(h)
        print(f"Priority {priority}: {name}")


if __name__ == "__main__":
    demo_basic_operations()
    demo_heapify()
    demo_peek()
    demo_tuple_heap()

#Topic: Max-Heap -- Negative Conversion Trick
# Time  : push O(log n), pop O(log n)
#Space: O(n) -- Heap size

import heapq


def max_heappush(h: list[int], val: int) -> None:
    """Adding an element to a Max-Heap -- convert to negative and push."""
    heapq.heappush(h, -val)


def max_heappop(h: list[int]) -> int:
    """Extracting the maximum value from a Max-Heap -- restore the sign after pop."""
    return -heapq.heappop(h)


def max_peek(h: list[int]) -> int:
    """Query the maximum value of a Max-Heap (without removal)."""
    return -h[0]


def demo_max_heap() -> None:
    """Max-Heap demo."""
    print("=" * 50)
    print("Max-Heap -- Negative Conversion Trick")
    print("=" * 50)

    h: list[int] = []
    values = [5, 3, 8, 1, 9, 2, 7]

    for val in values:
        max_heappush(h, val)
        print(f"push({val}) -> Maximum: {max_peek(h)}, Internal heap: {h}")

    print()
    print("Pop order (always maximum):")
    snapshot = list(h)
    while snapshot:
        val = max_heappop(snapshot)
        print(f"  pop() -> {val}")


def demo_max_heap_with_heapify() -> None:
    """Convert the existing list into a Max-Heap."""
    print()
    print("=" * 50)
    print("Max-Heap heapify")
    print("=" * 50)

    lst = [3, 1, 4, 1, 5, 9, 2, 6, 5]
    print("Original:", lst)

    #Convert all elements to negative and then heapify
    neg_lst = [-x for x in lst]
    heapq.heapify(neg_lst)
    print("Internal heap (negative):", neg_lst)
    print("Maximum:", -neg_lst[0])

    #Extract all elements (in descending order)
    sorted_desc = []
    while neg_lst:
        sorted_desc.append(-heapq.heappop(neg_lst))
    print("Descending sort:", sorted_desc)


def demo_kth_largest() -> None:
    """Find the K-th largest element -- Maintain a Min-Heap of size K."""
    print()
    print("=" * 50)
    print("k-th largest element -- Maintain min-heap of size K")
    print("=" * 50)

    def kth_largest(nums: list[int], k: int) -> int:
        """Time O(n log k), Space O(k)."""
        h: list[int] = []
        for num in nums:
            heapq.heappush(h, num)
            if len(h) > k:
                heapq.heappop(h)    #Remove the smallest element
        return h[0]   #Root of size k min-heap = k-th largest value

    test_cases = [
        ([3, 2, 1, 5, 6, 4], 2, 5),
        ([3, 2, 3, 1, 2, 4, 5, 5, 6], 4, 4),
    ]
    for nums, k, expected in test_cases:
        result = kth_largest(nums, k)
        ok = "OK" if result == expected else "NG"
        print(f"nums={nums}, k={k} -> {result}  {ok}")


if __name__ == "__main__":
    demo_max_heap()
    demo_max_heap_with_heapify()
    demo_kth_largest()

#Topic: Top-K Patterns -- nlargest, nsmallest, Streaming Top-K
#Time: O(n log k) -- processing n elements, maintaining a heap of size k
#Space: O(k) -- heap size

import heapq


def top_k_largest_builtin(nums: list[int], k: int) -> list[int]:
    """Using heapq.nlargest -- O(n log k)."""
    return heapq.nlargest(k, nums)


def top_k_smallest_builtin(nums: list[int], k: int) -> list[int]:
    """Using heapq.nsmallest -- O(n log k)."""
    return heapq.nsmallest(k, nums)


def top_k_largest_manual(nums: list[int], k: int) -> list[int]:
    """Maintain Top-K maximum values with a size k min-heap -- O(n log k).

    Idea: If the root of the min-heap < new element, replace the root.
    The k elements remaining in the heap at the end are the k overall maximum values."""
    h: list[int] = []
    for num in nums:
        if len(h) < k:
            heapq.heappush(h, num)
        elif num > h[0]:          #Replace if larger than the current k-th
            heapq.heapreplace(h, num)  #heappop + heappush (optimization)
    return sorted(h, reverse=True)


def streaming_top_k(stream: list[int], k: int) -> list[int]:
    """Maintaining Top-K in streaming data in real-time.

    Keep the Top-K each time an element comes in.
    Time O(n log k), Space O(k)."""
    h: list[int] = []
    results = []
    for num in stream:
        heapq.heappush(h, num)
        if len(h) > k:
            heapq.heappop(h)
        results.append(sorted(h, reverse=True))
    return results[-1] if results else []


def top_k_frequent(nums: list[int], k: int) -> list[int]:
    """Return top-K elements based on frequency -- O(n log k)."""
    freq: dict[int, int] = {}
    for num in nums:
        freq[num] = freq.get(num, 0) + 1

    #nlargest with (frequency, value)
    return [num for num, _ in heapq.nlargest(k, freq.items(), key=lambda x: x[1])]


if __name__ == "__main__":
    nums = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5]
    k = 3

    print("=" * 55)
    print(f"Top-K Test (nums={nums}, k={k})")
    print("=" * 55)

    print("nlargest (top 3):", top_k_largest_builtin(nums, k))
    print("nsmallest (bottom 3):", top_k_smallest_builtin(nums, k))
    print("manual top-k largest:", top_k_largest_manual(nums, k))

    print()
    print("=" * 55)
    print("Streaming Top-K")
    print("=" * 55)
    stream = [7, 2, 5, 1, 9, 3]
    k2 = 3
    final = streaming_top_k(stream, k2)
    print(f"stream={stream}, k={k2}")
    print("Final Top-K:", final)

    print()
    print("=" * 55)
    print("Top-K by frequency")
    print("=" * 55)
    freq_cases = [
        ([1, 1, 1, 2, 2, 3], 2, [1, 2]),
        ([1], 1, [1]),
    ]
    for arr, k3, expected in freq_cases:
        result = top_k_frequent(arr, k3)
        ok = "OK" if sorted(result) == sorted(expected) else "NG"
        print(f"nums={arr}, k={k3} -> {result}  {ok}")

    print()
    print("=" * 55)
    print("nlargest vs sorted complexity comparison guide")
    print("=" * 55)
    print("nlargest(k, arr): O(n log k) -- advantageous when k << n")
    print("sorted(arr)[-k:]: O(n log n) -- similar when k is close to n")
    print("However, if k >= n/2, sorted may be faster")

#Topic: Heap Sort -- O(n log n)
#Time: O(n log n) -- heapify O(n) + n times heappop O(log n)
#Space: O(n) heapq version / O(1) in-place version

import heapq


def heap_sort_asc(arr: list[int]) -> list[int]:
    """Heap sort ascending order -- using heapq, O(n log n).

    1. Convert list to a min-heap O(n)
    2. Repeat heappop n times -> ascending order O(n log n)"""
    h = list(arr)         #Copy (original preserved)
    heapq.heapify(h)      # O(n)
    result = []
    while h:
        result.append(heapq.heappop(h))   #O(log n) * n times
    return result


def heap_sort_desc(arr: list[int]) -> list[int]:
    """Heap sort descending order -- max-heap trick, O(n log n)."""
    h = [-x for x in arr]
    heapq.heapify(h)
    result = []
    while h:
        result.append(-heapq.heappop(h))
    return result


def heap_sort_inplace(arr: list[int]) -> None:
    """In-place heap sort -- Space O(1) (no additional array).

    Implement max heap directly for sorting.
    Python's heapq is a min-heap, so direct implementation is needed."""
    n = len(arr)

    def sift_down(arr: list[int], root: int, end: int) -> None:
        """Rearrange downward from the root position (max-heap)."""
        while True:
            largest = root
            left  = 2 * root + 1
            right = 2 * root + 2

            if left < end and arr[left] > arr[largest]:
                largest = left
            if right < end and arr[right] > arr[largest]:
                largest = right

            if largest == root:
                break
            arr[root], arr[largest] = arr[largest], arr[root]
            root = largest

    #1. Construct max-heap O(n)
    for i in range(n // 2 - 1, -1, -1):
        sift_down(arr, i, n)

    #2. Extract one by one and place at the end O(n log n)
    for end in range(n - 1, 0, -1):
        arr[0], arr[end] = arr[end], arr[0]  #Maximum value to the end
        sift_down(arr, 0, end)


if __name__ == "__main__":
    import random

    print("=" * 55)
    print("Heap Sort correctness test")
    print("=" * 55)

    test_cases = [
        [5, 3, 8, 1, 9, 2, 7, 4, 6],
        [1],
        [],
        [3, 3, 3],
        [9, 8, 7, 6, 5, 4, 3, 2, 1],
    ]

    for arr in test_cases:
        asc  = heap_sort_asc(arr)
        desc = heap_sort_desc(arr)

        #In-place verification
        inplace = list(arr)
        heap_sort_inplace(inplace)

        expected_asc  = sorted(arr)
        expected_desc = sorted(arr, reverse=True)

        ok_asc     = "OK" if asc     == expected_asc     else "NG"
        ok_desc    = "OK" if desc    == expected_desc    else "NG"
        ok_inplace = "OK" if inplace == expected_asc     else "NG"

        print(
            f"Input: {str(arr):30} |"
            f"Ascending order {ok_asc} |"
            f"Descending order {ok_desc} |"
            f"in-place {ok_inplace}"
        )

    print()
    print("=" * 55)
    print("Performance test (n=50,000)")
    print("=" * 55)
    import time

    big = random.sample(range(200_000), 50_000)

    t0 = time.perf_counter()
    heap_sort_asc(big)
    t1 = time.perf_counter()
    print(f"heap_sort_asc (heapq)  : {t1 - t0:.4f}s")

    t0 = time.perf_counter()
    sorted(big)
    t1 = time.perf_counter()
    print(f"sorted() (Timsort)     : {t1 - t0:.4f}s")
    print("(Timsort is faster built-in in CPython)")