Contains Duplicate II: дубликаты на расстоянии k, словарь или окно

Коротко

• Оба решения делают один проход и находят ответ на одном и том же шаге. Разница только в том, что они забывают.
• Словарь помнит каждое значение до конца прохода, множество помнит только последние $k$ элементов.
• При $n = 10^5$, $k = 10$ и массиве без повторов в замере ниже словарь занял почти 10 МБ, множество меньше 3 КБ.
• Если $k$ сравнимо с числом разных значений, выигрыш исчезает, а иногда множество проигрывает даже по памяти.

Что изменилось по сравнению с базовой задачей

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

Условие (на LeetCode это задача 219): дан массив nums и число k. Нужно понять, есть ли два разных индекса $i$ и $j$, для которых nums[i] == nums[j] и $|i - j| \le k$.

nums = [1, 2, 3, 1],       k = 3  →  true   (единицы на индексах 0 и 3, разница ровно 3)
nums = [1, 2, 3, 1, 2, 3], k = 2  →  false  (ближайшие равные отстоят на 3)

Перебор очевиден: для каждого элемента смотрим на следующие $k$ соседей. Это $O(n \cdot \min(n, k))$ по времени. При $n = k = 10^5$ получается порядка $10^{10}$ сравнений, так что дальше разговора он не уйдёт.

Слово «недавно» можно понять двумя способами, и из них получаются два решения:

1. Запоминать, когда значение встречалось в последний раз, и считать расстояние.
2. Держать под рукой только последние $k$ элементов. Всё, что старше, забывать: оно уже слишком далеко.

Решение 1: словарь последних индексов

def contains_nearby_duplicate(nums, k):
    last = {}
    for i, x in enumerate(nums):
        if x in last and i - last[x] <= k:
            return True
        last[x] = i
    return False


print(contains_nearby_duplicate([1, 2, 3, 1], 3))  # True
print(contains_nearby_duplicate([1, 2, 3, 1, 2, 3], 2))  # False

Почему хватает одного индекса на значение? Ближайшее слева вхождение $x$ — это и есть последнее. Если даже оно дальше $k$, то все более старые тем более. Поэтому индекс перезаписывается всегда, в том числе когда проверка не прошла.

Прогоним на [7, 3, 5, 7, 3, 3] с $k = 2$:

Время $O(n)$ в среднем, потому что операции со словарём в среднем стоят $O(1)$. Откуда берётся это «в среднем», написано в статье про хеш-таблицы. Память $O(m)$: в словаре по записи на каждое разное значение, и в худшем случае их $n$. Заметьте, что $k$ на память не влияет вообще.

Решение 2: множество как окно размера k

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

def contains_nearby_duplicate(nums, k):
    window = set()
    for i, x in enumerate(nums):
        if x in window:
            return True
        window.add(x)
        if len(window) > k:
            window.remove(nums[i - k])
    return False


print(contains_nearby_duplicate([1, 2, 3, 1], 3))  # True
print(contains_nearby_duplicate([1, 2, 3, 1, 2, 3], 2))  # False
print(contains_nearby_duplicate([1, 1], 0))  # False

Тот же пример, [7, 3, 5, 7, 3, 3] и $k = 2$:

Ответ получен на том же шаге, что и в словаре.

В общем случае окну нужен счётчик, а не множество, но тут это лишнее. Второй экземпляр значения внутри окна означает, что ответ найден и функция уже завершилась. Значит, все элементы окна различны, и remove убирает ровно один элемент, который точно там лежит.

Время $O(n)$: каждый элемент один раз добавляется и один раз удаляется. Память $O(\min(k, m))$ плюс один элемент на мгновение между add и remove.

Это один и тот же алгоритм

Возьмите словарь и начните удалять из него всё, что вышло за окно:

def contains_nearby_duplicate(nums, k):
    last = {}
    for i, x in enumerate(nums):
        if x in last:
            return True
        last[x] = i
        if i >= k:
            del last[nums[i - k]]
    return False


print(contains_nearby_duplicate([1, 2, 3, 1], 3))  # True
print(contains_nearby_duplicate([1, 2, 3, 1, 2, 3], 2))  # False

Теперь сам факт наличия ключа гарантирует, что расстояние подходит. Значения словаря больше никто не читает, и его можно заменить множеством. Вся разница между двумя решениями в том, выбрасываем ли мы устаревшее.

Сколько памяти это стоит на самом деле

Асимптотика обещает выигрыш, но цифры интереснее. Я прогнал оба решения на CPython 3.9.6 с $n = 10^5$: tracemalloc для пика памяти, время — лучшее из семи запусков. Ответ на всех наборах False, поэтому оба прошли массив до конца. На других версиях и машинах цифры будут другими, смотрите на соотношения.

Что отсюда следует:

• Выигрыш множества зависит от отношения $k$ к $m$. При малом $k$ разница достигает трёх порядков. При большом почти всё съедается.
• Множество может проиграть по памяти. В последней строке окно держит почти все разные значения, и множество заняло вдвое больше словаря. Удаления оставляют в его таблице «мёртвые» слоты, и она перестраивается с запасом. Это детали CPython, а не свойство алгоритма, но на собеседовании полезно понимать: $O(\min(k, m))$ не значит «всегда меньше».
• Словарь быстрее. В этом прогоне примерно в 1,5–2,5 раза: множество на каждом шаге делает лишнее удаление. Асимптотика у обоих $O(n)$, но если памяти хватает, платить за окно незачем.
• В Python у словаря есть скрытая цена. Индексы больше 256 — отдельные объекты по 28 байт. На 1000 значениях таблица занимает около 37 КБ (sys.getsizeof), и ещё примерно 28 КБ приходится на сами индексы. Отсюда и 66 КБ.

Что выбрать

По умолчанию — словарь. Шесть строк, меньше мест для ошибки на единицу, и он быстрее.

Множество — когда $k$ мало по сравнению с числом разных значений. Например, проверка уникальности событий за последние $k$ записей на длинной истории: память не растёт с длиной потока.

Словарь, если нужно что-то кроме да/нет. Частые продолжения задачи решаются им почти без изменений: вернуть пару индексов (return last[x], i) или найти минимальное расстояние между равными элементами. Окно индексов не хранит, поэтому ему эти вопросы даются хуже.

def min_gap(nums):
    last = {}
    best = None
    for i, x in enumerate(nums):
        if x in last:
            gap = i - last[x]
            best = gap if best is None else min(best, gap)
        last[x] = i
    return best


print(min_gap([7, 3, 5, 7, 3, 3]))  # 1
print(min_gap([1, 2, 3]))  # None

Что меняется для потока

Решение с окном выше опирается на nums[i - k], то есть на весь массив в памяти. Если данные приходят потоком, порядок элементов в окне придётся хранить самому:

from collections import deque


def stream_has_nearby_duplicate(stream, k):
    window = set()
    order = deque()
    for x in stream:
        if x in window:
            return True
        window.add(x)
        order.append(x)
        if len(window) > k:
            window.remove(order.popleft())
    return False


print(stream_has_nearby_duplicate(iter([1, 2, 3, 1]), 3))  # True
print(stream_has_nearby_duplicate(iter([1, 2, 3, 1]), 2))  # False

deque хранит только ссылки, в замере при $k = 1000$ он добавил около 8 КБ. А вот словарь без чистки на бесконечном потоке с постоянно новыми значениями рано или поздно съест всю память. Именно здесь разница между решениями перестаёт быть академической.

Типичные ошибки

Не обновлять индекс, если расстояние оказалось большим. Индекс записывают только в ветке else:

def wrong(nums, k):
    last = {}
    for i, x in enumerate(nums):
        if x in last:
            if i - last[x] <= k:
                return True
        else:
            last[x] = i  # индекс повторного вхождения теряется
    return False


print(wrong([1, 2, 1, 1], 1))  # False, а правильно True: индексы 2 и 3

Проверять индекс на истинность. if last.get(x) and ... выглядит короче, но индекс 0 превращается в False, поэтому [1, 1] с $k = 2$ даст False. В JavaScript та же ловушка с if (last[x]). Проверяйте наличие ключа.

Ошибка на единицу в размере окна. Если держать $k + 1$ элемент до проверки, то [1, 2, 1] с $k = 1$ вернёт True, хотя единицы отстоят на 2. Если удалить уходящий элемент до проверки, [1, 2, 3, 1] с $k = 3$ вернёт False, хотя расстояние ровно 3. Держите в голове инвариант: перед проверкой в окне лежат последние $k$ элементов, и прогоните пример, где расстояние равно $k$.

< вместо <=. Та же граница, но в словаре. [1, 1] с $k = 1$ должен вернуть True.

Отдельная ветка для $k = 0$. Она не нужна. Двух разных индексов на расстоянии 0 не бывает, и оба решения сами возвращают False.

Список вместо множества. x in window на списке стоит $O(k)$, и всё решение деградирует до перебора.

Усложнение: Contains Duplicate III

Возможное продолжение на интервью (на LeetCode это задача 220). Теперь значения не обязаны совпадать, они могут отличаться не больше чем на $t$:

$|i - j| \le k, \quad |nums[i] - nums[j]| \le t$

Окно индексов остаётся, а вот проверка x in window больше не работает: нужно найти в окне что-то близкое к $x$, и хеш-таблица такого запроса не понимает. Обратите внимание, что словарь «последний индекс» тут сразу выпадает: подойти может любое значение из диапазона, а перебрать их все нельзя. Поэтому разговор начинается с окна.

Отсортированное окно

Первая идея: держать окно отсортированным и смотреть на соседей $x$ по значению. Ближайшие к $x$ числа окна стоят сразу слева и справа от позиции вставки.

import bisect


def contains_almost_duplicate(nums, k, t):
    window = []
    for i, x in enumerate(nums):
        pos = bisect.bisect_left(window, x)
        if pos < len(window) and window[pos] - x <= t:
            return True
        if pos > 0 and x - window[pos - 1] <= t:
            return True
        window.insert(pos, x)
        if i >= k:
            window.pop(bisect.bisect_left(window, nums[i - k]))
    return False


print(contains_almost_duplicate([1, 2, 3, 1], 3, 0))  # True
print(contains_almost_duplicate([1, 5, 9, 1, 5, 9], 2, 3))  # False

На сбалансированном дереве это $O(n \log \min(n, k))$: TreeSet с floor и ceiling в Java, SortedList из пакета sortedcontainers в Python. Код выше использует обычный список: поиск логарифмический, зато insert и pop сдвигают хвост за $O(k)$. На практике это работает быстро, но интервьюеру стоит назвать честную оценку.

Корзины

Следующий вопрос напрашивается сам: можно ли за $O(n)$? Можно.

Разобьём числовую ось на корзины шириной $t + 1$. Для t = 3 это [0..3], [4..7], [8..11] и так далее. Номер корзины считается как x // (t + 1). Дальше работают два наблюдения:

• Два числа в одной корзине отличаются максимум на $t$. Значит, если корзина уже занята, ответ найден. Поэтому в каждой корзине никогда не лежит больше одного числа, и хватает словаря «корзина → значение».
• Подходящая пара может лежать только в соседней корзине. Числа через одну корзину отличаются минимум на $t + 2$. Проверяем b - 1 и b + 1, причём по-настоящему: там разница может оказаться и больше $t$.

def contains_almost_duplicate(nums, k, t):
    width = t + 1
    buckets = {}
    for i, x in enumerate(nums):
        b = x // width
        if b in buckets:
            return True
        if b - 1 in buckets and x - buckets[b - 1] <= t:
            return True
        if b + 1 in buckets and buckets[b + 1] - x <= t:
            return True
        buckets[b] = x
        if i >= k:
            del buckets[nums[i - k] // width]
    return False


print(contains_almost_duplicate([1, 2, 3, 1], 3, 0))  # True
print(contains_almost_duplicate([1, 5, 9, 1, 5, 9], 2, 3))  # False
print(contains_almost_duplicate([-3, 3], 1, 3))  # False

Прогон на [1, 5, 9, 1, 5, 9], $k = 2$, $t = 3$, ширина корзины 4:

Ответ False. Время $O(n)$, память $O(\min(n, k))$. Удалять корзину при выходе из окна безопасно: второго числа в ней нет, иначе функция уже вернула бы True.

При $t = 0$ ширина равна единице, номер корзины совпадает со значением, проверка соседей никогда не срабатывает. Получается в точности множество-окно из решения 2. Хороший способ проверить себя и заодно показать интервьюеру, что вы видите связь.

Где здесь ошибаются:

• Ширина ровно $t + 1$. С шириной $t + 2$ правило «одна корзина — ответ найден» перестаёт работать: [0, 4] при $k = 1$, $t = 3$ попадут в одну корзину и дадут True, хотя разница 4. А ширина $t$ при $t = 0$ — это деление на ноль.
• Отрицательные числа. В Python // округляет вниз, и всё работает: -3 // 4 даёт -1, а 3 // 4 даёт 0. В Java -3 / 4 и 3 / 4 оба дают 0, поэтому -3 и 3 попадают в одну корзину при $t = 3$ и получается ложное срабатывание. Нужен Math.floorDiv, в JavaScript — Math.floor.
• Переполнение. В языках с фиксированной шириной int считайте разности и t + 1 в long, если диапазон значений в условии это допускает.

Как рассказывать на собеседовании

Задача лёгкая, и код на ней много не покажет, поэтому очки набираются на разговоре:

1. Назовите перебор и его цену одной фразой.
2. Напишите словарь. Объясните, почему хватает последнего индекса.
3. Сами скажите, что память $O(m)$ не зависит от $k$, и предложите окно на случай малого $k$ или потока.
4. Назовите условие, при котором окно не окупается, и скажите, что словарь быстрее.
5. Перед сдачей прогоните границы: расстояние ровно $k$, $k = 0$ и $k$ больше длины массива.

Если спросят про разницу значений, начинайте с отсортированного окна и переходите к корзинам. Корзины редко придумывают с нуля за пять минут, а вот рассуждение «как получить проверку за $O(1)$, если точное совпадение заменить диапазоном» стоит проговорить вслух.

Что запомнить

• Словарь «последний индекс» хранит $O(m)$ и не зависит от $k$, множество-окно хранит $O(\min(k, m))$.
• Словарь с чисткой устаревших ключей и есть множество-окно.
• Окно окупается, когда $k$ намного меньше числа разных значений, и почти обязательно на потоке.
• Счётчик в окне не нужен: дубликат внутри окна сразу означает ответ.
• Contains Duplicate III решается корзинами шириной $t + 1$ за $O(n)$, а при $t = 0$ сводится к обычному окну.

Попробуйте оба решения в нашем тренажёре, а затем, для тренировки инварианта, напишите окно без подсказок: у него больше мест для ошибки на единицу, чем кажется на первый взгляд. Задачи с разницей значений в тренажёре пока нет, её можно решить на LeetCode.