Add median in a stream using heap-based approach with doctests

Author: Bharatgwl

data_structures/heap/median_in_a_stream.py

@@ -1,14 +1,16 @@
+"""
+Median in a stream using a heap-based approach.
+
+Reference:
+https://en.wikipedia.org/wiki/Median#Running_median
+"""
+
 import heapq
 
 
 def signum(a: int, b: int) -> int:
     """
-    Compare two integers.
-
-    Returns:
-        1 if a > b
-       -1 if a < b
-        0 if a == b
+    Return 1 if a > b, -1 if a < b, 0 if equal.
     """
     if a > b:
         return 1
@@ -24,54 +26,51 @@ def call_median(
     median: int,
 ) -> int:
     """
-    Insert an element into heaps and update the median.
+    Update heaps and median based on the new element.
+
+    Args:
+        element (int): new element in stream
+        max_heap (list[int]): max heap (as negative numbers)
+        min_heap (list[int]): min heap
+        median (int): current median
+
+    Returns:
+        int: updated median
     """
-    case = signum(len(max_heap), len(min_heap))
+    size_diff = signum(len(max_heap), len(min_heap))
 
-    if case == 0:
+    if size_diff == 0:
         if element > median:
             heapq.heappush(min_heap, element)
             median = min_heap[0]
         else:
             heapq.heappush(max_heap, -element)
             median = -max_heap[0]
-
-    elif case == 1:
+    elif size_diff == 1:
         if element > median:
             heapq.heappush(min_heap, element)
         else:
-            heapq.heappush(min_heap, -heapq.heappop(max_heap))
-            heapq.heappush(max_heap, -element)
+            heapq.heappush(min_heap, -heapq.heappushpop(max_heap, -element))
         median = (-max_heap[0] + min_heap[0]) // 2
-
-    else:
+    else:  # size_diff == -1
         if element > median:
-            heapq.heappush(max_heap, -heapq.heappop(min_heap))
-            heapq.heappush(min_heap, element)
+            heapq.heappush(max_heap, -heapq.heappushpop(min_heap, element))
         else:
             heapq.heappush(max_heap, -element)
         median = (-max_heap[0] + min_heap[0]) // 2
 
     return median
 
 
-def median_in_a_stream(numbers: list[int]) -> list[int]:
+def median_in_a_stream(arr: list[int]) -> list[int]:
     """
-    Find the median after each insertion in a stream of integers.
-
-    Uses two heaps and follows the classic running median logic.
-
-    Reference:
-    https://en.wikipedia.org/wiki/Median#Running_median
+    Return the median after each new element in the stream.
 
     Args:
-        numbers: List of integers
+        arr (list[int]): list of integers
 
     Returns:
-        List of medians after each insertion
-
-    Raises:
-        ValueError: If the input list is empty
+        list[int]: running medians
 
     >>> median_in_a_stream([20, 14, 13, 16, 17])
     [20, 17, 14, 15, 16]
@@ -82,16 +81,24 @@ def median_in_a_stream(numbers: list[int]) -> list[int]:
     ...
     ValueError: Input list must not be empty
     """
-    if not numbers:
+    if not arr:
         raise ValueError("Input list must not be empty")
 
-    max_heap: list[int] = []
-    min_heap: list[int] = []
-    median = 0
-    result: list[int] = []
+    max_heap: list[int] = []  # left side (as negative numbers)
+    min_heap: list[int] = []  # right side
+    median = arr[0]
+    max_heap.append(-arr[0])
+    medians: list[int] = [median]
 
-    for element in numbers:
+    for element in arr[1:]:
         median = call_median(element, max_heap, min_heap, median)
-        result.append(median)
+        medians.append(median)
+
+    return medians
+
 
-    return result
+if __name__ == "__main__":
+    n = int(input("Enter number of elements: ").strip())
+    arr = [int(input().strip()) for _ in range(n)]
+    result = median_in_a_stream(arr)
+    print("Running medians:", result)