658. Find K Closest Elements — разбор задачи

1. Условие задачи

Дано: отсортированный массив целых чисел arr, число k и число x.

Нужно: вернуть k чисел из массива, которые ближе всего к x. Ответ должен быть отсортирован по возрастанию.

Правило «кто ближе»: число a ближе к x, чем b, если |a - x| < |b - x|. А если расстояния равны — побеждает меньшее число.

Ограничения: массив не пустой, до 10^4 элементов, 1 <= k <= len(arr). Числа и x от -10^4 до 10^4. Важно: массив отсортирован неубывающе, то есть дубликаты разрешены — и именно на них прячутся грабли этой задачи. Ещё момент: сам x в массиве может отсутствовать.

Пример: arr = [1, 2, 3, 4, 5], k = 4, x = 3 → [1, 2, 3, 4]. Числа 2, 3, 4 очевидно ближайшие, а за четвёртое место борются 1 и 5: расстояния равны (по 2), побеждает меньшее — единица.

2. Решение

Ключевое наблюдение: массив отсортирован, поэтому k ближайших к x чисел — это всегда сплошное окно длины k. Не могут ближайшие числа стоять с «дыркой»: если элемент внутри дырки, он ближе к x, чем края. Значит, вся задача сводится к одному вопросу: где начинается это окно?

Решение в лоб (Brute Force)

Сортируем весь массив по расстоянию до x: sorted(arr, key=lambda num: abs(num - x)), берём первые k штук и сортируем их обратно по возрастанию. Работает за O(N log N) и вообще не использует то, что массив уже отсортирован. Годится как первый ответ на собеседовании, но интервьюер сразу спросит: «А можно быстрее?»

Линейное решение: сужаем окно с двух концов

Раз ответ — окно, начнём с самого большого окна (весь массив) и будем выкидывать по одному дальнему элементу, пока не останется ровно k. Кандидаты на вылет всегда только два: самый левый и самый правый элемент — они дальше всех от x. Сравнили расстояния — выкинули того, кто дальше. При ничьей выкидываем правого (он больше, а при равных расстояниях мы обязаны оставить меньших).

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        left, right = 0, len(arr) - 1
        # Пока в окне больше k элементов — выкидываем дальний край
        while right - left + 1 > k:
            if x - arr[left] > arr[right] - x:
                left += 1   # левый дальше — выкидываем левый
            else:
                right -= 1  # правый дальше или ничья — выкидываем правый
        return arr[left:right + 1]

Это O(N): в худшем случае делаем N - k шагов сужения. Уже хорошо, и на собеседовании такое решение обычно засчитывают. Но у задачи стоят теги Binary Search, и самое красивое решение — впереди.

Идея оптимального решения: бинпоиск по левой границе окна

Не будем искать сами элементы — будем искать индекс left, с которого начинается окно ответа arr[left : left + k]. Этот индекс живёт в диапазоне от 0 до len(arr) - k (правее окно уже не влезет). А по этому диапазону можно ходить бинпоиском.

Берём кандидата mid. Как понять, окно [mid, mid + k - 1] нас устраивает или его надо сдвинуть вправо? Сравниваем два элемента: arr[mid] (левый край текущего окна) и arr[mid + k] (первый элемент за правым краем). Ровно один из них лишний:

если x - arr[mid] > arr[mid + k] - x — левый край дальше от x, чем сосед справа за окном. Значит, выгодно сдвинуться вправо: left = mid + 1;
иначе (в том числе при ничьей) — окно [mid, mid + k - 1] нас устраивает: запоминаем кандидата ans = mid и пробуем найти начало ещё левее: right = mid - 1.

Когда указатели разошлись (left > right), в ans лежит самое левое годное начало окна.

Одна техническая деталь: кандидата mid = len(arr) - k проверять нельзя — для него сосед arr[mid + k] уже за пределами массива. Поэтому бинпоиск гоняем по right = len(arr) - k - 1, а в ans заранее кладём len(arr) - k: если все проверки скажут «вправо» — окно стоит у самого правого края.

Внимание, грабли: сравниваем именно x - arr[mid] > arr[mid + k] - x, со знаком, без abs(). Почему — в «Пояснениях» ниже.

Алгоритм по шагам

Кладём в ans = len(arr) - k самое правое возможное начало окна.
Ставим left = 0, right = len(arr) - k - 1 — кандидаты, у которых сосед arr[mid + k] существует.
Пока left <= right, берём mid = (left + right) // 2 и сравниваем x - arr[mid] и arr[mid + k] - x.
Если левый край дальше — окно правее: left = mid + 1.
Иначе — mid годится: запоминаем ans = mid и ищем левее: right = mid - 1.
Возвращаем срез arr[ans : ans + k].

Визуализация на `arr = [1, 2, 3, 4, 5]`, `k = 2`, `x = 4`

Старт: left = 0, right = 3 - 1 = 2, ans = 3.

Шаг	left	right	mid	x - arr[mid]	arr[mid+k] - x	Вывод	Действие	ans
1	0	2	1	4 - 2 = 2	4 - 4 = 0	левый край дальше	left = 2	3
2	2	2	2	4 - 3 = 1	5 - 4 = 1	ничья — вправо не идём	ans = 2, right = 1	2
3	2	1	—	—	—	left > right, стоп	ответ `arr[2:4] = [3, 4]`	2

Обрати внимание на шаг 2: расстояния равны, и мы не идём вправо — автоматически выполняем правило «при ничьей берём меньшие».

Код на Python

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        # 1. Запасной ответ — окно у самого правого края
        ans = len(arr) - k
        # 2. Кандидаты, для которых arr[mid + k] существует
        left, right = 0, len(arr) - k - 1
        while left <= right:
            mid = (left + right) // 2
            # 3. Сравниваем расстояния СО ЗНАКОМ, без abs!
            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1      # левый край дальше — окно правее
            else:
                ans = mid           # mid годится — запоминаем
                right = mid - 1     # и пробуем найти начало левее
        # 4. Ответ — окно из k элементов начиная с ans
        return arr[ans:ans + k]

Краевые случаи

k == len(arr) — тогда right = -1, цикл не запускается ни разу, ans = 0 — возвращаем весь массив. Красиво само собой работает.
x меньше всех элементов — условие «вправо» ни разу не сработает (слева расстояния «отрицательные»), ans доедет до нуля: первые k элементов.
x больше всех элементов — все проверки скажут «вправо», ans останется запасным len(arr) - k: последние k элементов.
x отсутствует в массиве — вообще не важно: мы ни разу не ищем сам x, мы ищем границу окна.
Дубликаты и равные расстояния — за счёт сравнения без abs окно корректно прижимается влево, к меньшим числам.

3. Сложность алгоритма

Бинпоиск идёт по диапазону из N - k + 1 возможных начал окна — O(log(N - k)).
Срез arr[left:left + k] для ответа — O(k).
Дополнительной памяти не используем — пара переменных-указателей.

Итог по времени: O(log(N - k) + k).

Итог по памяти: O(1) дополнительной (сам ответ на k элементов не считаем).

Пояснения: почему нельзя `abs()`

Соблазн велик: «сравниваем расстояния» — значит abs(x - arr[mid]) > abs(arr[mid + k] - x). На большинстве тестов это даже пройдёт, но ломается на дубликатах. Классический контрпример: arr = [1, 1, 2, 2, 2, 2, 2, 3, 3], x = 3, k = 3.

Правильный ответ — [2, 3, 3]: две тройки на расстоянии 0 и одна двойка на расстоянии 1. Версия с abs выдаёт [2, 2, 2]: при равных модулях расстояний она не понимает, в какую сторону двигаться, и уводит окно не туда. А сравнение со знаком знает направление: если оба элемента левее x, то x - arr[mid] большое положительное, а arr[mid + k] - x отрицательное — и мы уверенно шагаем вправо.

Есть красивая геометрическая интерпретация: условие x - arr[mid] > arr[mid + k] - x — это то же самое, что x > (arr[mid] + arr[mid + k]) / 2. То есть мы просто смотрим, по какую сторону от середины окна лежит x, и двигаем окно в ту сторону. А при ничьей (в том числе когда все элементы одинаковые и все расстояния нулевые) уходим в ветку ans = mid, right = mid - 1 — прижимаемся к меньшим числам, как и требует условие.

Пояснения: связь с шаблоном feasible-функции

Наш код — это в чистом виде универсальный шаблон «бинпоиск по монотонной булевой функции» (как он сформулирован, например, в Algo Monster): находим feasible-функцию, которая превращает индексы в булев массив вида FFF...TTT, и ищем первый True.

Здесь feasible-функция от кандидата mid — «окно с началом в mid не выгодно сдвигать вправо»:

def feasible(mid: int) -> bool:
    return x - arr[mid] <= arr[mid + k] - x

Она монотонна: с ростом mid левая часть x - arr[mid] не растёт, а правая arr[mid + k] - x не убывает — значит, если feasible стала истинной, правее она истинна всегда. Получаем ровно картинку FFF...TTT, а ответ задачи — первый True (самое левое годное начало окна). Дальше шаблон применяется механически: feasible(mid) истинна → ans = mid, right = mid - 1, иначе left = mid + 1.

Единственное отличие от каноничного шаблона: там стартовое значение ответа -1 («вдруг True вообще нет»). У нас ответ есть всегда: последний кандидат len(arr) - k feasible «по определению» — правее сдвигаться просто некуда (и проверить его формулой нельзя: arr[mid + k] за массивом). Поэтому вместо -1 кладём в ans именно его — это гарантированный True в конце булева массива. В яндексовской версии за O(log k) такую же роль играет min(idx, len(nums) - k).

Вариация от Яндекса

На собеседованиях в Яндексе дают такую версию:

Дан массив nums, отсортированный в неубывающем порядке, индекс idx и число k. Нужно найти k ближайших к значению nums[idx] чисел в массиве и вернуть в любом порядке. При равных расстояниях предпочтение отдаётся меньшим числам.

Чем отличается от оригинала:

Вместо произвольного x дают индекс — целимся в x = nums[idx], и это значение гарантированно есть в массиве.
Ответ можно вернуть в любом порядке — жизнь стала только проще.
Правило ничьей то же самое: меньшие числа побеждают.

Главный подарок: нам выдали позицию цели бесплатно. В оригинале, чтобы найти место x в массиве, нужен бинпоиск — здесь он уже не нужен. Поэтому самое естественное решение — расширять окно от idx двумя указателями: начинаем с окна из одного элемента nums[idx] и k - 1 раз забираем более близкого соседа — слева или справа. При ничьей берём левого (он меньше).

def find_k_closest(nums: list[int], idx: int, k: int) -> list[int]:
    x = nums[idx]
    left = right = idx  # окно из одного элемента — nums[idx]
    # 1. Расширяем окно k - 1 раз
    for _ in range(k - 1):
        if left == 0:
            right += 1          # слева стенка — берём справа
        elif right == len(nums) - 1:
            left -= 1           # справа стенка — берём слева
        elif x - nums[left - 1] <= nums[right + 1] - x:
            left -= 1           # левый ближе или ничья — берём меньший
        else:
            right += 1          # правый строго ближе
    return nums[left:right + 1]

Сложность: O(k) по времени, O(1) по памяти — даже без бинпоиска, потому что искать точку старта не нужно. Заметь зеркальную симметрию с линейным решением оригинала: там мы сужали окно, выкидывая дальних (при ничьей выкидываем правого), здесь расширяем, забирая ближних (при ничьей берём левого). Правило одно и то же — меньшие числа в приоритете.

Бинпоиск границ за `O(log k)`

Бинпоиск из оригинала работает здесь дословно (x = nums[idx], дальше тот же код), но знание idx даёт ему апгрейд. В оригинале мы гоняли поиск по всему диапазону начал окна [0, N - k] — отсюда O(log(N - k)). Но раз цель — сам элемент массива, окно ответа обязано накрывать значение nums[idx] (у него расстояние 0 — оно точно в ответе). Значит, начало окна лежит в отрезке [idx - k + 1, idx] (с подрезкой под границы массива) — а в нём всего k кандидатов. Бинпоиск по нему — O(log k) вместо O(log(N - k)).

def find_k_closest(nums: list[int], idx: int, k: int) -> list[int]:
    x = nums[idx]
    # 1. Начало окна — где-то в пределах k позиций от idx
    ans = min(idx, len(nums) - k)     # запасной ответ — самый правый кандидат
    left, right = max(0, idx - k + 1), ans - 1
    # 2. Тот же бинпоиск, но по диапазону длины <= k
    while left <= right:
        mid = (left + right) // 2
        if x - nums[mid] > nums[mid + k] - x:
            left = mid + 1
        else:
            ans = mid
            right = mid - 1
    return nums[ans:ans + k]

Два тонких момента, которые стоит проговорить интервьюеру:

Почему правая граница idx корректна: условие «сдвинуться вправо» при mid = idx требует 0 > nums[idx + k] - x, а это невозможно: справа от idx все элементы >= x. То есть правее idx окно не уедет никогда.
При длинных сериях дубликатов x суженный поиск может найти другое начало окна, чем поиск по полному диапазону (другие индексы внутри серии равных), но набор значений в окне будет тем же самым — а вернуть нас просят именно числа, да ещё и в любом порядке.

Честная оговорка про итоговую сложность: сами границы мы теперь находим за O(log k), но если надо вернуть сами k чисел — на вывод всё равно уйдёт O(k), и суммарно получится O(k) — как и у двух указателей. Так что выигрыш чистого бинпоиска проявляется, когда достаточно отдать границы окна (индексы), а не копировать элементы. На собеседовании рассказать оба способа и эту оговорку — беспроигрышный ход.

Вывод

Паттерн задачи: ответ — сплошное окно фиксированной длины в отсортированном массиве, а бинпоиск можно запускать не по элементам, а по позициям начала окна. Плюс важный урок про сравнение расстояний со знаком вместо abs — именно на этом задача валит тех, кто решает «на автомате».

🐆