Coding Interview Prep · Урок

Медиана потока данных и слияние k списков

Поддерживайте две кучи — максимальную для меньшей половины и минимальную для большей — чтобы обновлять медиану за O(log n), а затем объединяйте k отсортированных списков с помощью кучи

Урок 4 из 413 шагов

«Медиана потока данных и слияние k списков» — бесплатный урок Coding Interview Prep на CoddyKit. Это урок 4 из 4. Ты можешь прочитать весь урок бесплатно ниже — а потом практиковать его прямо в браузере с встроенным редактором кода и ИИ-репетитором 24/7. Это часть пути обучения Coding Interview Prep, и твой прогресс синхронизируется между веб-версией и приложением CoddyKit. Курс Coding Interview Prep содержит 4 уроков всего.

Задача о медиане потока данных

Найдите медиану потока данных (LeetCode #295) — задача, в которой нужно эффективно поддерживать две операции: addNum(num) для add числа и findMedian() для возврата текущей медианы. Медиана списка чётной длины — это среднее двух центральных значений. При использовании отсортированного списка прямолинейное решение даёт вставку за O(n) и получение медианы за O(1). Оптимальное решение использует две кучи для вставки за O(log n) и получения медианы за O(1).

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')

Реализация 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

Понимание того, почему поддерживается инвариант двух куч, крайне важно для объяснения решения на собеседовании. Рассмотрим пошаговое добавление [5, 15, 1, 3]. После каждой вставки выполните balance, чтобы нижняя макс-куча содержала меньшую половину элементов. Инвариант гарантирует, что max(lo) <= min(hi) всегда выполняется, поэтому медиана доступна непосредственно на вершине одной или обеих куч.

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})')

Медиана в скользящем окне

Медиана в скользящем окне (LeetCode #480) — более сложный вариант задачи: найдите медиану каждого окна размера k по мере его перемещения по массиву. Подход с двумя кучами расширяется с помощью множества отложенного удаления, чтобы обрабатывать элементы, выходящие из окна. Когда элемент покидает окно, пометьте его в множестве удалений; когда он окажется на вершине любой кучи, удалите его.

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-стороннем слиянии

merge k отсортированных списков (LeetCode #23) — фундаментальная задача, применяемая во внешней сортировке, слиянии баз данных и распределённых системах. Даны k отсортированных связных списков, содержащих в общей сложности n узлов; объедините их в один отсортированный список. Наивный подход — выполнять merge по два списка за раз — имеет сложность O(kn), а при использовании метода «разделяй и властвуй» — O(n log k). Подход с кучей обрабатывает каждый узел ровно один раз, выполняя O(log k) работы на узел, то есть имеет общую сложность O(n log k).

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-стороннее слияние с мин-кучей

Инициализируйте кучу первым узлом каждого списка. На каждом шаге выполните pop минимального элемента, add его в результат и push следующий узел из того же списка, если он есть. В куче всегда не более k элементов — по одному началу каждого активного списка. Поскольку всего обрабатывается n узлов, а для каждого выполняются операции с кучей за O(log k), общая временная сложность составляет O(n log k), а пространственная — O(k) для кучи.

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 списков

Наименьший диапазон (LeetCode #632) — это задача поиска наименьшего диапазона [lo, hi], в котором находится хотя бы один элемент из каждого из 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-й наименьший элемент в матрице

k-й наименьший элемент в отсортированной матрице (LeetCode #378): дана матрица n×n, в которой отсортированы каждая строка и каждый столбец. Найдите k-й наименьший элемент. Рассматривайте каждую строку как отсортированный список и используйте k-стороннее слияние с кучей. Другой вариант — выполнить двоичный поиск по диапазону значений. Подход с кучей работает за O(k log n), что эффективно при небольшом k; двоичный поиск работает за O(n log(max-min)) и лучше подходит для больших 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

Две кучи для текущей статистики

Шаблон с двумя кучами применим не только к медиане. С его помощью можно поддерживать текущий quantile, например 25-й перцентиль: настройте размер нижней кучи так, чтобы в ней находилось p*n элементов, а верхней — (1-p)*n элементов. При каждом добавлении элемента восстанавливайте баланс, как раньше. Этот шаблон встречается в задачах потоковой статистики, где одновременно нужны эффективная вставка и запросы quantile.

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 ближайших к началу координат точек

k ближайших к началу координат точек (LeetCode #973) — задача, в которой используется макс-куча размера k. Для каждой точки выполните push её квадрат расстояния, чтобы не вычислять квадратный корень. Когда размер кучи превысит k, выполните pop самой удалённой точки. Оставшиеся k точек являются k ближайшими. Сложность составляет O(n log k). Другой вариант — алгоритм быстрого выбора со средней сложностью O(n), но решение с кучей проще правильно реализовать и объяснить на собеседовании.

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}')

Две кучи: анализ времени и памяти

Подход с двумя кучами для медианы обеспечивает O(log n) на каждую операцию addNum и O(1) на findMedian. Для хранения всех элементов требуется O(n) памяти. k-стороннее слияние занимает O(n log k) времени и O(k) памяти для кучи. Эти оценки близки к оптимальным: можно доказать нижнюю границу Omega(n log k) в сравнительной модели для k-стороннего слияния, показывающую, что решение с кучей асимптотически оптимально. На собеседованиях всегда чётко называйте эти сложности.

# 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)')

Быстрая проверка

Проверьте, насколько хорошо Вы поняли концепции «Структуры данных и алгоритмы — подготовка к собеседованию по программированию» из этого урока.

Итоги урока

В этом уроке Вы изучили: двухкучечный MedianFinder со вставкой за O(log n) и получением медианы за O(1), k-стороннее слияние с мин-кучей за O(n log k) времени и O(k) памяти, а также расширения, включая медиану в скользящем окне, наименьший диапазон и k ближайших точек. Далее мы рассмотрим представление графов и настройку их обхода.

Можно начать бесплатно

Изучай Coding Interview Prep с ИИ-репетитором — бесплатно

Пиши и запускай код прямо в браузере, получай мгновенную помощь от ИИ-репетитора 24/7 и продолжи учиться на сайте или в приложении.

Курсы
90
Уроки
360

Часто задаваемые вопросы

Урок «Медиана потока данных и слияние k списков» бесплатный?

Да — полный текст урока «Медиана потока данных и слияние k списков» бесплатно доступен здесь в веб-версии. Чтобы практиковать его интерактивно (встроенный редактор кода и ИИ-репетитор 24/7) и разблокировать остальной курс Coding Interview Prep, подпишись на CoddyKit PRO. Курс Coding Interview Prep содержит 4 уроков всего.

Чему я научусь в уроке «Медиана потока данных и слияние k списков»?

Поддерживайте две кучи — максимальную для меньшей половины и минимальную для большей — чтобы обновлять медиану за O(log n), а затем объединяйте k отсортированных списков с помощью кучи Ты практикуешь Coding Interview Prep с помощью реального кода, который запускаешь прямо в браузере, и ИИ-репетитор 24/7 отвечает на твои вопросы во время урока.

Нужен ли мне опыт, чтобы начать Coding Interview Prep?

Предыдущий опыт не требуется. Coding Interview Prep на CoddyKit структурирован для всех уровней — от новичков до продвинутых, поэтому ты можешь начать отсюда или с самого начала и учиться в своем темпе. Это урок 4 из 4.

Сколько времени занимает урок «Медиана потока данных и слияние k списков»?

Большинство уроков CoddyKit занимают около 5–10 минут. Каждый из них компактный и интерактивный, поэтому ты постоянно делаешь прогресс и продолжаешь с того же места в веб-версии и приложении.

Можно ли писать и запускать код в этом уроке Coding Interview Prep?

Да. Каждый урок Coding Interview Prep включает встроенный редактор кода, поэтому ты пишешь и запускаешь реальный код прямо в браузере и получаешь моментальную обратную связь от AI — локальная установка не требуется.

Все уроки этого курса

  1. Свойство кучи и представление массивом
  2. Построение кучи, добавление и извлечение с нуля
  3. heapq в Python и приёмы работы с максимальной кучей
  4. Медиана потока данных и слияние k списков
← Назад к Coding Interview Prep