K Closest Points to Origin

Q: Какова временная сложность алгоритма поиска K ближайших к началу координат точек?

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

Q: Почему для поиска k ближайших точек используют макс-кучу, а не мин-кучу?

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

Q: Тебе нужен квадратный корень, чтобы сравнивать расстояния?

Нет. Для неотрицательных чисел a < b выполняется тогда и только тогда, когда √a < √b , поэтому x² + y² задаёт тот же порядок, что и расстояние. Кроме того, это значение остаётся целым числом: здесь оно не превышает 2 × 10^8 , что укладывается в диапазон 32-битных чисел.

Q: Можно ли решить задачу «K ближайших к началу координат точек» за время O(n)?

В среднем — да. Quickselect распределяет точки относительно случайного опорного элемента в том же порядке (x² + y², x, y) и продолжает работу только с той частью, в которой находится позиция k , за O(n) в среднем. Ответ — это k точек, стоящих перед этой позицией; их сортировка добавляет O(k log k) .

СреднеКуча Сортировка

+10

На плоскости даны n точек: точка i имеет координаты (xs[i], ys[i]). Верните k точек, ближайших к началу координат (0, 0), каждую в виде пары [x, y].

Упорядочьте ответ по расстоянию от начала координат: сначала ближайшие точки. Если расстояния одинаковы, сначала укажите точку с меньшим x, затем — с меньшим y. Точка, указанная несколько раз, учитывается отдельно для каждого вхождения.

Функция

kClosest(xs: integer-array, ys: integer-array, k: integer) → integer-2d-array

xsinteger-array: координата x каждой точки
ysinteger-array: координата y каждой точки
kinteger: сколько баллов вернуть
Возвращаетinteger-2d-array: k ближайших точек в виде пар [x, y], сначала ближайшие; при равенстве расстояний — сначала по x, затем по y

Ограничения

n == xs.length == ys.length
1 ≤ k ≤ n ≤ 5000
-10⁴ ≤ xs[i], ys[i] ≤ 10⁴
Баллы могут повторяться; каждый экземпляр считается отдельным баллом.

Примеры

Ввод: xs = [3, -1, 2, 0]ys = [1, 2, -2, 5]k = 2
Вывод: [[-1, 2], [2, -2]]
Пояснение: Квадраты расстояний равны 10, 5, 8 и 25. Два наименьших значения соответствуют (-1, 2) со значением 5 и (2, -2) со значением 8, именно в таком порядке.

Ввод: xs = [1, -1, 0, 1]ys = [1, 1, 2, -1]k = 3
Вывод: [[-1, 1], [1, -1], [1, 1]]
Пояснение: (1, 1), (-1, 1) и (1, -1) имеют квадрат расстояния 2, а (0, 2) — 4. Среди точек с одинаковым значением у (-1, 1) наименьшая координата x, а (1, -1) стоит перед (1, 1) по координате y.

Ввод: xs = [7]ys = [-7]k = 1
Вывод: [[7, -7]]
Пояснение: Есть только одна точка, и вам нужна одна, поэтому ответ — сама точка (7, -7).

Дополнительный вопрос

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

Подсказки

Открывайте по одной. Каждая подсказывает чуть больше.

Открыто 0 из 3

Решение

Для расстояния не нужен квадратный корень: x² + y² задаёт тот же порядок и остаётся целым числом. Вам нужны только k из n точек, поэтому max-куча из k ближайших на данный момент точек, где самая дальняя находится наверху и готова к замене, требует O(n log k) вместо полной сортировки.

Выберите ближайшую точку k раз

Верно, но не успевает на самых больших тестах

Идея

Первая точка в ответе — ближайшая из всех, вторая — ближайшая из оставшихся, и так далее. Поэтому просканируйте точки k раз, каждый раз выбирая среди ещё не выбранных точку, которая идёт первой по (x² + y², x, y).

При каждом сканировании проверяются все n точек. При n = 5000 и k = 2500 это 1.25 × 10^7 трёхкомпонентных сравнений.

Алгоритм

Отметь каждую точку как не выбранную.
Повтори k раз: просматривай не выбранные точки и найди ту, которая идет первой по (x² + y², x, y).
Отметь ее как выбранную и добавь [x, y] к ответу.
Верни ответ.

def kClosest(xs, ys, k):
    n = len(xs)
    taken = [False] * n
    result = []
    for _ in range(k):
        best = -1
        for i in range(n):
            # Compare (squared distance, x, y): the order the answer uses.
            if not taken[i] and (best == -1 or
                                 (xs[i] * xs[i] + ys[i] * ys[i], xs[i], ys[i]) <
                                 (xs[best] * xs[best] + ys[best] * ys[best], xs[best], ys[best])):
                best = i
        taken[best] = True
        result.append([xs[best], ys[best]])
    return result

Упорядочьте все точки

Идея

Ответ — это первые k точек в полном порядке, поэтому отсортируй каждую точку по (x² + y², x, y) и оставь первые k. В примере 1 отсортированные ключи — (5, -1, 2), (8, 2, -2), (10, 3, 1) и (25, 0, 5), и первые две точки дают ответ.

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

Алгоритм

Сформируй ключ (x² + y², x, y) для каждой точки.
Отсортируй точки по ключу.
Верни первые k из них в виде пар [x, y].

def kClosest(xs, ys, k):
    points = list(zip(xs, ys))
    # Sort by squared distance, then x, then y: the order the answer uses.
    points.sort(key=lambda p: (p[0] * p[0] + p[1] * p[1], p[0], p[1]))
    return [[x, y] for x, y in points[:k]]

Оставьте k ближайших элементов в max-heap

Идея

Храните лучшие точки k, встреченные к этому моменту. Новая точка может вытеснить только сохранённую точку, занимающую последнее место в рейтинге, поэтому храните сохранённые точки в max-heap по ключу (x² + y², x, y), поместив такую точку на вершину. Когда сохранено k точек, точка, занимающая более высокое место в рейтинге, заменяет вершину; любую другую точку отбрасывают, потому что k уже сохранённых точек занимают более высокие места.

В примере 2 при k = 3 после первых трёх точек на вершине находится (0, 2) с ключом (4, 0, 2). Точка (1, -1) имеет ключ (2, 1, -1), занимает более высокое место в рейтинге и заменяет её.

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

Алгоритм

Создайте пустую max-кучу с ключом (x² + y², x, y).
Вставьте каждую из первых k точек.
Для каждой последующей точки сравните её с верхним элементом: если она имеет более высокий приоритет, замените им верхний элемент и просейте его вниз; в противном случае пропустите её.
Удалите верхний элемент k раз; каждая удалённая точка заполняет последнюю свободную позицию в ответе.
Верните ответ.

import heapq


def kClosest(xs, ys, k):
    # heapq is a min-heap; negating every part of (distance, x, y) puts the
    # farthest kept point, in the answer's order, on top.
    heap = []
    for x, y in zip(xs, ys):
        item = (-(x * x + y * y), -x, -y)
        if len(heap) < k:
            heapq.heappush(heap, item)
        elif item > heap[0]:
            # The new point comes before the farthest kept one: swap them.
            heapq.heapreplace(heap, item)
    # Popping gives the farthest first, so fill the answer from the back.
    result = [None] * len(heap)
    for i in range(len(result) - 1, -1, -1):
        _, x, y = heapq.heappop(heap)
        result[i] = [-x, -y]
    return result

Ловушки и крайние случаи

Пропуск разрешения ничьей. В примере 2 три точки находятся на расстоянии 2; без сравнения сначала x, а затем y они будут расположены в порядке ввода, а не как (-1, 1), (1, -1), (1, 1).
Хранение k точек в мин-куче. Её вершина — ближайшая точка, поэтому при каждом удалении отбрасывается точка, которая нужна в ответе, и в итоге остаются k самых дальних точек.
Получение ответа напрямую из массива кучи. Этот массив лишь частично упорядочен; сначала извлеките точки по одной.
Удаление повторяющихся точек. Если (5, 5) указана три раза и k = 3, в ответе эта точка будет три раза.

Частые вопросы4

Какова временная сложность алгоритма поиска K ближайших к началу координат точек?

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

Почему для поиска k ближайших точек используют макс-кучу, а не мин-кучу?

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

Тебе нужен квадратный корень, чтобы сравнивать расстояния?

Нет. Для неотрицательных чисел a < b выполняется тогда и только тогда, когда √a < √b, поэтому x² + y² задаёт тот же порядок, что и расстояние. Кроме того, это значение остаётся целым числом: здесь оно не превышает 2 × 10^8, что укладывается в диапазон 32-битных чисел.

Можно ли решить задачу «K ближайших к началу координат точек» за время O(n)?

В среднем — да. Quickselect распределяет точки относительно случайного опорного элемента в том же порядке (x² + y², x, y) и продолжает работу только с той частью, в которой находится позиция k, за O(n) в среднем. Ответ — это k точек, стоящих перед этой позицией; их сортировка добавляет O(k log k).

def kClosest(xs, ys, k):
    # Напишите код здесь

Тест #1

Тест #2

Тест #3

Аргументы

xs = [3, -1, 2, 0]
ys = [1, 2, -2, 5]
k = 2

Ожидаемый вывод

[[-1, 2], [2, -2]]

K Closest Points to Origin

Функция

Ограничения

Примеры

Подсказки

Решение

Выберите ближайшую точку k раз

Упорядочьте все точки

Оставьте k ближайших элементов в max-heap

Ловушки и крайние случаи

Частые вопросы4

Похожие задачи