K ближайших точек

Постановка задачи

sqrt(x^2 + y^2)

points = [[1, 3], [-2, 2]]
k = 1

[[-2, 2]]

points = [[3, 3], [5, -1], [-2, 4]]
k = 2

[[3, 3], [-2, 4]]

Решение

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

Идея решения:

Будем поддерживать максимальную кучу размера k. В ней хранятся текущие k ближайших точек. Вершина кучи - самая дальняя точка среди выбранных кандидатов.

Когда мы берем очередную точку, добавляем ее в кучу. Если размер кучи стал больше k, удаляем вершину (самую дальнюю точку). После этого в куче снова остаются k ближайших точек из всех просмотренных.

Нам не нужно вычислять квадратный корень sqrt(x^2 + y^2). Для сравнения расстояний достаточно использовать квадрат расстояния x^2 + y^2. Если одно расстояние меньше другого, то и квадрат этого расстояния тоже меньше.

Основные шаги:

1. Создаем максимальную кучу по квадрату расстояния.
2. Проходим по всем точкам.
3. Добавляем текущую точку в кучу.
4. Если размер кучи больше k, удаляем самую дальнюю точку.
5. Возвращаем все точки, оставшиеся в куче.

Рассмотрим пример:

points = [[3, 3], [5, -1], [-2, 4]]
k = 2

Квадраты расстояний:

[3, 3]  -> 18
[5, -1] -> 26
[-2, 4] -> 20

1. Добавляем [3, 3], куча хранит одну точку.

2. Добавляем [5, -1], куча хранит две точки.

3. Добавляем [-2, 4], размер стал 3. Удаляем самую дальнюю точку [5, -1] с квадратом расстояния 26.

В куче остались [3, 3] и [-2, 4].

Оценка сложности

Временная сложность

Мы обрабатываем n точек. Для каждой точки выполняем добавление в кучу и, возможно, удаление. Размер кучи не превышает k, поэтому каждая операция занимает O(log k).

Итоговая временная сложность: O(n log k).

Пространственная сложность

Куча хранит не более k точек, поэтому дополнительная память равна O(k).

Код решения

Приведем код решения.

import "github.com/emirpasic/gods/trees/binaryheap"
 
type Point struct {
    x    int
    y    int
    dist int
}
 
func kClosest(points [][]int, k int) [][]int {
    heap := binaryheap.NewWith(func(a, b interface{}) int {
        return b.(Point).dist - a.(Point).dist
    })
 
    for _, point := range points {
        x, y := point[0], point[1]
        heap.Push(Point{x: x, y: y, dist: x*x + y*y})
        if heap.Size() > k {
            heap.Pop()
        }
    }
 
    result := [][]int{}
    for !heap.Empty() {
        val, _ := heap.Pop()
        point := val.(Point)
        result = append(result, []int{point.x, point.y})
    }
    return result
}

Итоги

Задача K Closest Points to Origin показывает прием с максимальной кучей: чтобы найти k минимальных элементов, удобно поддерживать максимальную кучу размера k.

1. Суть алгоритма: добавляем точки в максимальную кучу и удаляем самую дальнюю, если кандидатов стало больше k.
2. Эффективность:
   • Временная сложность: O(n log k).
   • Пространственная сложность: O(k).
3. Граничные случаи:
   • k = points.length: в ответ попадут все точки.
   • Отрицательные координаты: квадрат расстояния x^2 + y^2 всегда неотрицателен, поэтому знаки координат не влияют на корректность вычислений.
4. Оптимальное решение: мы не сортируем все точки. Куча размера k хранит только ближайших кандидатов и сразу отбрасывает лишние элементы.