データストリームの中央値とK-wayマージ

データストリームの中央値問題

import heapq

# Strategy: maintain two halves of the data
# max_heap: lower half (stores negated values for max behavior)
# min_heap: upper half
# Invariant: len(max_heap) == len(min_heap) or len(max_heap) == len(min_heap) + 1
# Invariant: max(max_heap) <= min(min_heap)
# Median:
#   odd count:  max_heap[0] (top of lower half)
#   even count: average of tops of both halves
print('Two-heap strategy for O(log n) insert, O(1) median')

2つのヒープによるMedianFinderの実装

import heapq

class MedianFinder:
    def __init__(self):
        self.lo = []  # max-heap (negated) for lower half
        self.hi = []  # min-heap for upper half

    def addNum(self, num):
        heapq.heappush(self.lo, -num)   # push to lower half
        # Ensure max of lower <= min of upper
        if self.hi and -self.lo[0] > self.hi[0]:
            heapq.heappush(self.hi, -heapq.heappop(self.lo))
        # Balance sizes: lo can have at most 1 more than hi
        if len(self.lo) > len(self.hi) + 1:
            heapq.heappush(self.hi, -heapq.heappop(self.lo))
        elif len(self.hi) > len(self.lo):
            heapq.heappush(self.lo, -heapq.heappop(self.hi))

    def findMedian(self):
        if len(self.lo) > len(self.hi):
            return -self.lo[0]  # odd count: top of lower half
        return (-self.lo[0] + self.hi[0]) / 2

mf = MedianFinder()
for n in [1, 2, 3, 4, 5]: mf.addNum(n)
print(mf.findMedian())  # 3.0

MedianFinderの手順を追う

import heapq

# Manual trace for [5, 15, 1, 3]:
# add 5:   lo=[-5]        hi=[]       median=5
# add 15:  lo=[-5]        hi=[15]     median=(5+15)/2=10
# add 1:   lo=[-5,-1]     hi=[15]     median=5
# add 3:   lo=[-5,-3,-1]  hi=[15]     -- lo too big
#       -> lo=[-5,-3]      hi=[1,15]  -- wait, wrong direction
# Actually:
# add 1:   push to lo -> lo=[-5,-1], then 1>lo? No, -lo[0]=5>15? No
#          lo has 2, hi has 1: balance -> move lo top to hi
#          lo=[-1], hi=[5,15]
# Median = (-lo[0] + hi[0])/2 = (1+5)/2 = 3
mf2 = MedianFinder()
for n, expected in [(5, 5.0), (15, 10.0), (1, 5.0), (3, 4.0)]:
    mf2.addNum(n)
    print(f'After adding {n}: median={mf2.findMedian()} (expected ~{expected})')

スライディングウィンドウの中央値

import heapq

def median_sliding_window(nums, k):
    lo = []  # max-heap (negated)
    hi = []  # min-heap
    removed = {}
    result = []

    def balance():
        # Move valid tops to correct side
        while lo and removed.get(-lo[0], 0) > 0:
            removed[-lo[0]] -= 1; heapq.heappop(lo)
        while hi and removed.get(hi[0], 0) > 0:
            removed[hi[0]] -= 1; heapq.heappop(hi)

    for i, num in enumerate(nums):
        heapq.heappush(lo, -num)
        heapq.heappush(hi, -heapq.heappop(lo))
        if len(hi) > len(lo): heapq.heappush(lo, -heapq.heappop(hi))
        if i >= k:
            out = nums[i - k]
            removed[out] = removed.get(out, 0) + 1
        balance()
        if len(lo) > len(hi): heapq.heappush(hi, -heapq.heappop(lo))
        if i >= k - 1:
            if len(lo) > len(hi): result.append(float(-lo[0]))
            else: result.append((-lo[0] + hi[0]) / 2.0)
    return result

print(median_sliding_window([1,3,-1,-3,5,3,6,7], 3))  # [1,-1,-1,3,5,6]

K-wayマージ：問題

class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next

# Build a linked list from a Python list
def build_list(arr):
    dummy = ListNode(0)
    curr = dummy
    for val in arr:
        curr.next = ListNode(val)
        curr = curr.next
    return dummy.next

# Convert linked list to Python list for printing
def to_list(head):
    result = []
    while head:
        result.append(head.val)
        head = head.next
    return result

print('K-way merge: O(n log k) using a min-heap of k heads')

最小ヒープによるK-wayマージ

import heapq

def merge_k_lists(lists):
    dummy = ListNode(0)
    curr = dummy
    heap = []
    for i, node in enumerate(lists):
        if node:
            heapq.heappush(heap, (node.val, i, node))
    while heap:
        val, i, node = heapq.heappop(heap)
        curr.next = node
        curr = curr.next
        if node.next:
            heapq.heappush(heap, (node.next.val, i, node.next))
    return dummy.next

lists = [
    build_list([1, 4, 5]),
    build_list([1, 3, 4]),
    build_list([2, 6])
]
result = merge_k_lists(lists)
print(to_list(result))  # [1, 1, 2, 3, 4, 4, 5, 6]

K個のリストを包含する最小範囲

import heapq

def smallest_range(nums):
    heap = []
    current_max = float('-inf')
    for i, lst in enumerate(nums):
        heapq.heappush(heap, (lst[0], i, 0))
        current_max = max(current_max, lst[0])
    best = [float('-inf'), float('inf')]
    while heap:
        current_min, list_idx, elem_idx = heapq.heappop(heap)
        if current_max - current_min < best[1] - best[0]:
            best = [current_min, current_max]
        if elem_idx + 1 >= len(nums[list_idx]):
            break  # one list exhausted
        next_val = nums[list_idx][elem_idx + 1]
        heapq.heappush(heap, (next_val, list_idx, elem_idx + 1))
        current_max = max(current_max, next_val)
    return best

print(smallest_range([[4,10,15,24,26],[0,9,12,20],[5,18,22,30]]))
# [20, 24]

行列中のk番目に小さい要素

import heapq

def kth_smallest_matrix(matrix, k):
    n = len(matrix)
    heap = [(matrix[0][0], 0, 0)]
    count = 0
    visited = {(0, 0)}
    while heap:
        val, r, c = heapq.heappop(heap)
        count += 1
        if count == k:
            return val
        # Push right neighbor
        if c + 1 < n and (r, c+1) not in visited:
            heapq.heappush(heap, (matrix[r][c+1], r, c+1))
            visited.add((r, c+1))
        # Push bottom neighbor
        if r + 1 < n and (r+1, c) not in visited:
            heapq.heappush(heap, (matrix[r+1][c], r+1, c))
            visited.add((r+1, c))
    return -1

matrix = [[1,5,9],[10,11,13],[12,13,15]]
print(kth_smallest_matrix(matrix, 8))  # 13

実行統計のための2つのヒープ

import heapq

# Generalised two-heap for arbitrary quantile p
# lo contains floor(p * count) elements
# hi contains the remaining elements
class QuantileFinder:
    def __init__(self, p):
        self.p = p  # quantile (e.g., 0.5 for median)
        self.lo = []  # max-heap
        self.hi = []  # min-heap
        self.count = 0

    def add(self, num):
        self.count += 1
        heapq.heappush(self.lo, -num)
        heapq.heappush(self.hi, -heapq.heappop(self.lo))
        # Target: lo should have floor(p * count) elements
        target_lo = int(self.p * self.count)
        while len(self.lo) < target_lo:
            heapq.heappush(self.lo, -heapq.heappop(self.hi))
        while len(self.lo) > target_lo:
            heapq.heappush(self.hi, -heapq.heappop(self.lo))

    def quantile(self):
        return -self.lo[0] if self.lo else self.hi[0]

qf = QuantileFinder(0.5)  # median
for n in [1, 2, 3, 4, 5, 6]: qf.add(n)
print(qf.quantile())  # 3 (median of 1-6)

原点に最も近いk個の点を見つける

import heapq

def k_closest(points, k):
    heap = []  # max-heap via negation
    for x, y in points:
        dist_sq = x*x + y*y
        heapq.heappush(heap, (-dist_sq, x, y))
        if len(heap) > k:
            heapq.heappop(heap)  # remove farthest
    return [[x, y] for _, x, y in heap]

points = [[1,3], [-2,2], [5,8], [0,1], [-1,-1]]
print(k_closest(points, 2))
# Two closest to origin: [0,1] (dist=1) and [-1,-1] (dist=2)

# Verify by distances:
for x, y in points:
    print(f'({x},{y}): dist^2 = {x*x+y*y}')

2つのヒープ：時間計算量と空間計算量の分析

# Complexity summary for heap applications:
# Problem               | Time per op  | Space
# ----------------------|--------------|------
# MedianFinder.addNum   | O(log n)     | O(n)
# MedianFinder.find     | O(1)         | -
# Merge k sorted lists  | O(n log k)   | O(k)
# Kth smallest matrix   | O(k log n)   | O(n)
# K closest points      | O(n log k)   | O(k)
# Task scheduler        | O(n log 26)  | O(26)
# Kth largest stream    | O(log k)     | O(k)
# Sliding window median | O(n log k)   | O(k)

print('Heap problems: identify k (heap size) vs n (input size)')

理解度チェック

レッスンのまとめ