K ближайших точек к началу координат

K ближайших точек к началу координат

Задача поиска K ближайших точек — это классическая задача на знание различных подходов: от наивного сортирования до оптимизации через приоритетные очереди. Это хорошо проверяет понимание алгоритмов и структур данных.

Понимание задачи

Нам нужно найти K точек, которые находятся ближе всего к началу координат (0, 0). Расстояние вычисляется по евклидовой формуле, но мы можем оптимизировать, используя квадрат расстояния (без sqrt).

Точки: [[1, 3], [-2, 2], [5, 8], [0, 1]]
K = 2

Расстояния:
[1, 3]:   sqrt(1 + 9) = sqrt(10) ≈ 3.16
[-2, 2]:  sqrt(4 + 4) = sqrt(8) ≈ 2.83  ← 2-я ближайшая
[5, 8]:   sqrt(25 + 64) = sqrt(89) ≈ 9.43
[0, 1]:   sqrt(0 + 1) = sqrt(1) = 1     ← 1-я ближайшая

Ответ: [[0, 1], [-2, 2]]

Подход 1: Сортировка O(n log n)

Просто отсортируем все точки по расстоянию и возьмём первые K.

import math

def kClosest_Sort(points, k):
    # Вычисляем расстояния и сортируем
    points.sort(key=lambda p: p[0]**2 + p[1]**2)
    return points[:k]

print(kClosest_Sort([[1, 3], [-2, 2], [5, 8], [0, 1]], 2))
# [[0, 1], [-2, 2]]

Сложность:

• Временная: O(n log n) — сортировка
• Пространственная: O(log n) — для рекурсии сортировки

Плюсы:

• Очень просто реализуется
• Хороший выбор когда K близко к N

Минусы:

• Не оптимально когда K намного меньше N

Подход 2: Приоритетная очередь (MaxHeap) O(n log k) ⭐

Это оптимальное решение для случая, когда K намного меньше N. Используем макс-куч (max heap) размером K.

import heapq

def kClosest_Heap(points, k):
    # Python heapq — это min heap, поэтому используем отрицательные значения
    heap = []
    
    for point in points:
        dist = point[0]**2 + point[1]**2
        # Кладём (-dist, point) для max heap
        if len(heap) < k:
            heapq.heappush(heap, (-dist, point))
        elif dist < -heap[0][0]:  # Если точка ближе, чем дальнейшая в куче
            heapq.heapreplace(heap, (-dist, point))
    
    # Извлекаем все точки из кучи
    return [point for _, point in heap]

print(kClosest_Heap([[1, 3], [-2, 2], [5, 8], [0, 1]], 2))
# Может вернуться в другом порядке: [[0, 1], [-2, 2]]

Сложность:

• Временная: O(n log k) — для каждой точки делаем операцию с кучей размера K
• Пространственная: O(k) — куча содержит K элементов

Плюсы:

• Отлично когда K << N
• Можно использовать для потоковых данных

Минусы:

• Более сложная реализация

Подход 3: Быстрый выбор (QuickSelect) O(n) ⭐⭐

QuickSelect — это оптимальное решение в среднем случае. Основан на алгоритме быстрой сортировки, но останавливается раньше.

import random

def kClosest_QuickSelect(points, k):
    def distance_squared(point):
        return point[0]**2 + point[1]**2
    
    def select(left, right, k_smallest):
        # Если осталось только одна точка
        if left == right:
            return
        
        # Выбираем случайный pivot
        pivot_index = random.randint(left, right)
        pivot_index = partition(left, right, pivot_index)
        
        # Если нашли K-й элемент
        if k_smallest == pivot_index:
            return
        elif k_smallest < pivot_index:
            select(left, pivot_index - 1, k_smallest)
        else:
            select(pivot_index + 1, right, k_smallest)
    
    def partition(left, right, pivot_index):
        pivot_dist = distance_squared(points[pivot_index])
        # Переносим pivot в конец
        points[pivot_index], points[right] = points[right], points[pivot_index]
        store_index = left
        
        # Партиционируем
        for i in range(left, right):
            if distance_squared(points[i]) < pivot_dist:
                points[store_index], points[i] = points[i], points[store_index]
                store_index += 1
        
        # Возвращаем pivot на правильное место
        points[right], points[store_index] = points[store_index], points[right]
        return store_index
    
    select(0, len(points) - 1, k - 1)
    return points[:k]

print(kClosest_QuickSelect([[1, 3], [-2, 2], [5, 8], [0, 1]], 2))
# [[0, 1], [-2, 2]]

Сложность:

• Средняя временная: O(n) — в среднем случае
• Худшая временная: O(n²) — если всегда выбираем плохой pivot
• Пространственная: O(1) — in-place

Плюсы:

• Очень быстро в среднем случае
• In-place операция
• Линейное время

Минусы:

• Модифицирует входной массив
• Худший случай O(n²)

Сравнение всех подходов

Практическая реализация с обработкой ошибок

def kClosest_Production(points, k):
    """
    Найти K ближайших точек к началу координат.
    
    Args:
        points: List[List[int]] - массив координат [x, y]
        k: int - количество точек
    
    Returns:
        List[List[int]] - K ближайших точек
    
    Raises:
        ValueError: если K > количества точек
    """
    if k <= 0:
        return []
    
    if k > len(points):
        raise ValueError(f"K={k} не может быть больше количества точек {len(points)}")
    
    if k == len(points):
        return points
    
    # Используем MaxHeap для оптимальности
    heap = []
    
    for x, y in points:
        dist_sq = x * x + y * y
        
        if len(heap) < k:
            # Отрицательное значение для max heap
            import heapq
            heapq.heappush(heap, (-dist_sq, x, y))
        elif dist_sq < -heap[0][0]:
            heapq.heapreplace(heap, (-dist_sq, x, y))
    
    return [[x, y] for _, x, y in heap]

Тестирование

def test_k_closest():
    test_cases = [
        ([[1, 3], [-2, 2], [5, 8], [0, 1]], 2, {(-2, 2), (0, 1)}),
        ([[1, 1]], 1, {(1, 1)}),
        ([[0, 0], [1, 1]], 2, {(0, 0), (1, 1)}),
        ([[3, 3], [5, -1], [-2, 4]], 2, {(3, 3), (-2, 4)}),
    ]
    
    for points, k, expected in test_cases:
        result = kClosest_Heap(points, k)
        result_set = set(tuple(p) for p in result)
        assert result_set == expected, f"Failed for {points}, k={k}"
    
    print("All tests passed!")

test_k_closest()

Оптимизация памяти

Если нужно сэкономить память и работаем с потоками данных:

def kClosest_Stream(point_stream, k):
    """
    Найти K ближайших из потока точек.
    Используется когда все точки не влезают в памяти.
    """
    import heapq
    
    heap = []
    
    for point in point_stream:
        x, y = point
        dist_sq = x * x + y * y
        
        if len(heap) < k:
            heapq.heappush(heap, (-dist_sq, point))
        elif dist_sq < -heap[0][0]:
            heapq.heapreplace(heap, (-dist_sq, point))
    
    return [point for _, point in heap]

Решение с визуализацией

def kClosest_Visualized(points, k):
    # Вычисляем расстояния
    distances = [(p[0]**2 + p[1]**2, p) for p in points]
    
    # Показываем вычисления
    print("Точки и их расстояния до (0, 0):")
    for dist_sq, point in sorted(distances):
        import math
        dist = math.sqrt(dist_sq)
        print(f"{point}: √{dist_sq} ≈ {dist:.2f}")
    
    print(f"\nK ближайших точек (K={k}):")
    result = sorted(distances)[:k]
    for dist_sq, point in result:
        print(f"{point}: √{dist_sq}")
    
    return [p for _, p in result]

Ключевые выводы

✅ Выбор алгоритма зависит от K и N:

• K близко к N → сортировка
• K намного меньше N → приоритетная очередь
• Нужна максимальная скорость → quick select

✅ Оптимизация: используй квадрат расстояния, чтобы избежать sqrt

✅ Для потоков данных: используй max heap с фиксированным размером K

✅ Production код: добавляй валидацию и обработку ошибок