Subsets: все подмножества массива тремя способами и Subsets II с дубликатами

Коротко

Subsets — одна из тех задач, которые на собеседовании решают за десять минут, если знаешь шаблон, и за сорок, если изобретаешь его на ходу. Ценна она тем, что на ней видны сразу три разных взгляда на перебор: дерево решений, наращивание ответа и двоичная запись числа. Общая теория перебора с возвратом есть в статье про бэктрекинг, здесь разберём одну задачу и её вариант с повторами так, чтобы шаблон потом переносился на Combination Sum, перестановки и разбиения.

Условие

Дан массив nums из различных целых чисел. Верните все возможные подмножества. В ответе не должно быть одинаковых подмножеств, порядок не важен.

Вход:  [1, 2, 3]
Выход: [[], [1], [2], [1, 2], [3], [1, 3], [2, 3], [1, 2, 3]]

На LeetCode это задача 78, ограничения маленькие: до 10 элементов. Это сразу подсказка: раз $n \le 10$, авторы ждут экспоненциального решения, и никакого хитрого трюка, который сделает его полиномиальным, нет. Ответ сам по себе экспоненциального размера.

Почему подмножеств ровно 2^n

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

Для [1, 2, 3] получается 8 подмножеств, для 10 элементов — 1024, для 20 — уже больше миллиона. Средний размер подмножества — $n/2$, поэтому суммарно в ответе $n \cdot 2^{n-1}$ чисел. Отсюда нижняя граница: любой алгоритм, который выписывает ответ целиком, работает не быстрее $O(n \cdot 2^n)$. Все три способа ниже её достигают, и спорить на собеседовании стоит не об асимптотике, а о том, какой код проще и понятнее.

Эта же мысль «два решения на элемент» и есть первый способ.

Способ 1. Бэктрекинг «брать или не брать»

Нарисуем дерево решений. На уровне i решаем судьбу nums[i]: левая ветка — взять, правая — не брать. Листья дерева на глубине $n$ — это и есть подмножества.

                       []
            взять 1 /       \ не брать 1
                [1]           []
              /     \       /    \
          [1,2]     [1]   [2]     []
          /  \     /  \   / \    /  \
     [1,2,3][1,2][1,3][1][2,3][2][3] []

Код повторяет рисунок один в один:

def subsets(nums):
    result = []
    path = []

    def dfs(i):
        if i == len(nums):
            result.append(path[:])
            return
        path.append(nums[i])  # берём nums[i]
        dfs(i + 1)
        path.pop()  # откатываем выбор
        dfs(i + 1)  # не берём nums[i]

    dfs(0)
    return result


print(subsets([1, 2, 3]))
# [[1, 2, 3], [1, 2], [1, 3], [1], [2, 3], [2], [3], []]

Здесь три детали, на которых чаще всего ошибаются.

Копия path[:]. В result нужно класть снимок текущего пути, а не сам список. Если написать result.append(path), все элементы result окажутся ссылками на один и тот же объект, который к концу работы опустеет:

# result.append(path) вместо result.append(path[:])
# [[], [], [], [], [], [], [], []]

Восемь пустых списков — классический симптом. Если видите такое, ищите место, где сохраняется изменяемый объект без копии.

pop после рекурсивного вызова. Ветка «взять» добавила элемент, после возврата его надо убрать, иначе ветка «не брать» начнёт с испорченным состоянием. Это и есть «возврат» в слове «перебор с возвратом».

Глубина рекурсии. Она равна $n$, а при $n \le 20$ это ничтожно мало, так что о лимите рекурсии Python здесь можно не думать. Подробнее о том, как устроен стек вызовов, — в статье про рекурсию.

Вариант с циклом

Есть второй вид того же бэктрекинга, который встречается в разборах чаще, потому что он напрямую обобщается на Combination Sum и Subsets II. Вместо бинарного выбора на каждом шаге решаем: какой элемент добавить следующим, причём только из тех, что правее последнего взятого.

def subsets(nums):
    result = []
    path = []

    def dfs(start):
        result.append(path[:])  # каждый узел дерева — ответ
        for i in range(start, len(nums)):
            path.append(nums[i])
            dfs(i + 1)
            path.pop()

    dfs(0)
    return result


print(subsets([1, 2, 3]))
# [[], [1], [1, 2], [1, 2, 3], [1, 3], [2], [2, 3], [3]]

Разница в форме дерева. В варианте «брать/не брать» ответы лежат только в листьях, а всего узлов $2^{n+1} - 1$. Здесь ответом является каждый узел, и узлов ровно $2^n$. Условие i + 1 в рекурсивном вызове гарантирует, что элементы идут по возрастанию индексов, поэтому [2, 1] никогда не появится рядом с [1, 2].

Какой из двух вариантов писать — дело вкуса. Но если следующим вопросом будет Subsets II, удобнее сразу начать со второго.

Способ 2. Итеративное удвоение

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

def subsets(nums):
    result = [[]]
    for x in nums:
        result += [subset + [x] for subset in result]
    return result


print(subsets([1, 2, 3]))
# [[], [1], [2], [1, 2], [3], [1, 3], [2, 3], [1, 2, 3]]

Пошагово для [1, 2, 3]:

Размер списка удваивается на каждом шаге — отсюда название. Рекурсии нет, кода три строки, и его легко объяснить вслух. Минус один: все промежуточные подмножества хранятся сразу, так что ленивую генерацию по одному так не сделать. Для задачи «верни список» это не важно.

Тонкость Python: генератор списка [subset + [x] for subset in result] вычисляется целиком до того, как += начнёт расширять result. Поэтому цикл не уходит в бесконечность. Если переписать то же самое через for subset in result: result.append(...), получится бесконечный цикл: итерация пойдёт по только что добавленным элементам. Итерируйтесь по копии или по диапазону range(len(result)), зафиксированному заранее.

Способ 3. Битовые маски

Третий взгляд самый короткий. Запишем решение «брать или не брать» для всех $n$ элементов как строку из $n$ битов: бит i равен 1, если nums[i] взят. Каждая такая строка — двоичная запись числа от $0$ до $2^n - 1$, и наоборот, каждое число в этом диапазоне задаёт ровно одно подмножество.

Перебираем числа и для каждого проверяем биты:

def subsets(nums):
    n = len(nums)
    result = []
    for mask in range(1 << n):
        subset = [nums[i] for i in range(n) if mask >> i & 1]
        result.append(subset)
    return result


print(subsets([1, 2, 3]))
# [[], [1], [2], [1, 2], [3], [1, 3], [2, 3], [1, 2, 3]]

1 << n — это $2^n$, а mask >> i & 1 достаёт i-й бит. В Python приоритет сдвига выше, чем у &, поэтому скобки не нужны, но в других языках их лучше поставить явно. Вот то же на JavaScript:

function subsets(nums) {
  const n = nums.length
  const result = []
  for (let mask = 0; mask < 1 << n; mask++) {
    const subset = []
    for (let i = 0; i < n; i++) {
      if ((mask >> i) & 1) subset.push(nums[i])
    }
    result.push(subset)
  }
  return result
}

console.log(JSON.stringify(subsets([1, 2, 3])))
// [[],[1],[2],[1,2],[3],[1,3],[2,3],[1,2,3]]

Ограничение у способа одно: маска должна помещаться в число. В JavaScript побитовые операции работают с 32-битными целыми, так что 1 << 31 уже отрицательное, а 1 << 32 равно 1. Но подмножеств 30-элементного массива больше миллиарда, до этой границы вы не доберётесь по времени. В Python целые не ограничены.

Главная польза масок не в этой задаче, а в том, что подмножество становится обычным числом. Его можно положить в словарь, использовать как индекс массива, перебирать подмаски через sub = (sub - 1) & mask. На этом построено динамическое программирование по маскам, а базовые операции с битами разобраны в отдельной статье.

А можно через itertools?

Можно, и в рабочем коде так и стоит делать:

from itertools import combinations


def subsets(nums):
    return [list(c) for k in range(len(nums) + 1) for c in combinations(nums, k)]


print(subsets([1, 2, 3]))
# [[], [1], [2], [3], [1, 2], [1, 3], [2, 3], [1, 2, 3]]

На собеседовании такое решение обычно принимают как первый ответ, а потом просят написать без библиотеки. Полезно знать, что модуль есть, и уметь сказать, что он делает внутри. Подробнее о нём — в статье про itertools.

Сравнение способов

Обратите внимание: по асимптотике они одинаковы. Если интервьюер спрашивает «а можно быстрее?», правильный ответ — «нельзя, сам ответ имеет размер $n \cdot 2^n$». Это показывает, что вы понимаете, откуда берётся оценка, а не просто заучили её. Как вообще рассуждать о таких границах, есть в статье про O-нотацию.

Subsets II: когда в массиве есть повторы

Теперь условие меняется: элементы могут повторяться, а подмножества в ответе по-прежнему должны быть уникальными (LeetCode 90).

Вход:  [1, 2, 2]
Выход: [[], [1], [1, 2], [1, 2, 2], [2], [2, 2]]

Если применить любой способ выше без изменений, получится 8 подмножеств, и [2] с [1, 2] встретятся дважды: один раз с первой двойкой, другой раз со второй. Для нас это одно и то же подмножество.

Сколько ответов теперь

Пусть число v встречается c раз. Решение про него уже не бинарное: можно взять 0, 1, …, c копий, то есть $c + 1$ вариант. Всего уникальных подмножеств

Для [1, 2, 2] это $(1 + 1)(2 + 1) = 6$, что совпадает с ответом выше. Формула полезна для проверки: если ваш код выдаёт другое число, где-то дубли или потери.

Лобовое решение через set

Самое очевидное — сгенерировать всё и выкинуть повторы. Чтобы [1, 2] и [2, 1] не считались разными, массив сначала сортируют, а подмножества превращают в кортежи:

def subsets_with_dup(nums):
    seen = set()
    n = len(nums)
    nums = sorted(nums)
    for mask in range(1 << n):
        seen.add(tuple(nums[i] for i in range(n) if mask >> i & 1))
    return [list(t) for t in seen]

Работает, но тратит время на генерацию всех $2^n$ вариантов, даже если уникальных среди них мало. Для [2] * 10 это 1024 подмножества ради 11 уникальных. На собеседовании это решение стоит назвать и сразу предложить лучше.

Бэктрекинг с пропуском повторов

Идея: отсортировать массив, чтобы одинаковые числа стояли рядом, и запретить на одном уровне дерева выбирать одно и то же значение дважды.

def subsets_with_dup(nums):
    nums = sorted(nums)
    result = []
    path = []

    def dfs(start):
        result.append(path[:])
        for i in range(start, len(nums)):
            if i > start and nums[i] == nums[i - 1]:
                continue  # это значение на этом уровне уже пробовали
            path.append(nums[i])
            dfs(i + 1)
            path.pop()

    dfs(0)
    return result


print(subsets_with_dup([1, 2, 2]))
# [[], [1], [1, 2], [1, 2, 2], [2], [2, 2]]

Здесь всё держится на условии i > start. Разберём, что оно значит. Цикл внутри dfs(start) выбирает, каким будет следующий элемент подмножества. Если на этой позиции мы уже пробовали двойку с индексом 1, то пробовать двойку с индексом 2 бессмысленно: всё, что из неё вырастет, уже выросло из первой. А вот если i == start, то этот элемент первый кандидат на текущем уровне, и пропускать его нельзя — даже если он равен предыдущему. Именно так получается [2, 2]: на первом уровне берём nums[1], на втором уровне start = 2, и nums[2] уже не повтор, а законный первый кандидат.

Самая частая ошибка — написать i > 0 вместо i > start:

# if i > 0 and nums[i] == nums[i - 1]: continue
print(subsets_with_dup_wrong([1, 2, 2]))
# [[], [1], [1, 2], [2]]

Пропали [1, 2, 2] и [2, 2]: вторая двойка теперь отбрасывается всегда, в том числе когда её нужно взять вслед за первой. Проверка по формуле сразу это ловит — 4 вместо 6.

Вторая частая ошибка — забыть отсортировать. Тогда для [2, 1, 2] двойки не стоят рядом, сравнение с соседом их не видит, и дубли возвращаются.

Тот же приём в варианте «брать/не брать»

В бинарном дереве повторы убираются по-другому: если мы решили не брать значение, то не берём и все его копии подряд. Иначе подмножество «без первой двойки, но со второй» совпадёт с «с первой двойкой, без второй».

def subsets_with_dup(nums):
    nums = sorted(nums)
    result = []
    path = []

    def dfs(i):
        if i == len(nums):
            result.append(path[:])
            return
        path.append(nums[i])
        dfs(i + 1)
        path.pop()
        j = i
        while j < len(nums) and nums[j] == nums[i]:
            j += 1
        dfs(j)  # не берём nums[i] и все его копии

    dfs(0)
    return result


print(subsets_with_dup([1, 2, 2]))
# [[1, 2, 2], [1, 2], [1], [2, 2], [2], []]

Это та же формула $c + 1$, только в коде: взять 1, 2, …, c копий получается через цепочку «взять», а вариант «0 копий» — через прыжок dfs(j).

Итеративное удвоение с повторами

Удвоение тоже можно починить. Когда приходит число, равное предыдущему, добавлять его ко всем подмножествам нельзя: к «старым» его уже добавляли на прошлом шаге, получатся дубли. Добавлять нужно только к тем подмножествам, которые появились на прошлом шаге.

def subsets_with_dup(nums):
    nums = sorted(nums)
    result = [[]]
    start = 0
    for i, x in enumerate(nums):
        if i > 0 and x == nums[i - 1]:
            begin = start  # только новые с прошлого шага
        else:
            begin = 0  # все
        start = len(result)
        result += [subset + [x] for subset in result[begin:]]
    return result


print(subsets_with_dup([1, 2, 2]))
# [[], [1], [2], [1, 2], [2, 2], [1, 2, 2]]

Для [1, 2, 2]: после единицы [[], [1]]. Первая двойка добавляется ко всем: появились [2], [1, 2], start = 2. Вторая двойка совпадает с предыдущей, поэтому добавляется только к result[2:] — к [2] и [1, 2]. Итог: 6 подмножеств.

Битовые маски с повторами работают плохо: маска различает позиции, а не значения, так что без set или дополнительных проверок дубли не убрать. Это хороший ответ на вопрос «а почему не маски?».

Сложность Subsets II

В худшем случае, когда все элементы различны, это обычный Subsets: $O(n \cdot 2^n)$. Плюс $O(n \log n)$ на сортировку, что на фоне экспоненты незаметно. Когда повторов много, оптимальные решения работают пропорционально числу уникальных подмножеств, а лобовое через set всё равно перебирает все $2^n$. В этом и разница.

Как это связано с Combination Sum

Вариант с циклом for i in range(start, len(nums)) — это скелет целого семейства задач. Меняются три вещи: когда сохранять ответ, куда идти дальше и когда обрывать ветку.

В Combination Sum каждое число можно брать сколько угодно раз, поэтому рекурсия вызывается с тем же i, а не с i + 1. Combination Sum II — это буквально Subsets II, в котором сохраняются только подмножества с нужной суммой. Подробный разбор — в статье про Combination Sum, а решить задачу можно в тренажёре.

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

1. result.append(path) без копии. В ответе $2^n$ пустых списков.
2. Нет path.pop() после рекурсии. Подмножества «растут» и содержат лишние элементы.
3. dfs(start + 1) вместо dfs(i + 1). Пропадает часть ответов или появляются перестановки одного подмножества. Следующий старт зависит от того, какой элемент взяли, а не от того, где начался цикл.
4. Subsets II без сортировки. Сравнение с соседом не находит повторов, которые стоят не рядом.
5. i > 0 вместо i > start. Теряются подмножества с несколькими копиями одного числа.
6. Изменение списка во время итерации по нему в итеративном варианте. Бесконечный цикл.
7. Попытка ускорить до полинома. Ответ экспоненциальный, быстрее его выписать нельзя.

Как это спрашивают на собеседовании

Subsets — задача уровня Medium, но фактически разминочная. Типичный сценарий выглядит так.

Сначала просят решить базовую версию. Здесь важно назвать число ответов $2^n$ и объяснить, откуда оно, а потом написать любой из трёх способов без ошибок с копированием и откатом.

Потом спрашивают про сложность. Правильный ответ — $O(n \cdot 2^n)$ по времени, и полезно сразу добавить, что это нижняя граница из-за размера ответа.

Затем — follow-up с повторами. Если базовая версия была написана циклом от start, переход к Subsets II занимает две строки: сортировка и continue. Если бинарным выбором — нужен прыжок через копии. Интервьюер почти наверняка спросит, почему i > start, а не i > 0; ответ «чтобы не терять подмножества вроде [2, 2]» с примером закрывает вопрос.

Иногда добавляют: «а как сгенерировать подмножества без рекурсии?» (удвоение или маски) или «а если нужны подмножества ровно размера k?» (обрыв ветки по длине, это уже Combinations). Ещё один вариант — «верните подмножества лениво, по одному»: тогда бэктрекинг переписывается генератором с yield, а маски работают как есть, потому что каждое подмножество вычисляется из числа независимо от остальных.

Если хотите закрепить шаблон, решите подряд Subsets, Subsets II, Combinations, Combination Sum и Permutations. После пятой задачи код начнёт писаться сам, и на собеседовании останется время на разговор о деталях.