Давайте рассмотрим довольно классическую задачку на программирование под названием «Сумма трех чисел» (и производную от нее — «Сумму четырех чисел»). Решать будем брутфорс-методом, а затем усовершенствуем решение при помощи рекурсии так, чтобы оно подходило для любой задачи «Сумма K чисел». Вы увидите, что хотя брутфорс-решение не очень хорошо масштабируется, оно все равно не бесполезно.

Давайте начнем с простого варианта задачи — суммы трех чисел. Вот условие:

# Дан массив nums, содержащий n целых чисел,  # и дано целевое число - target. # Напишите функцию, которая будет определять,  # есть ли в nums три элемента a, b и c,  # сумма которых равна целевому числу # (т. е. a + b + c = target).  # Пример:  # Input: # nums = [-1,0,1,2,-1,-4] # target = 0  # Output: [[-1,-1,2],[-1,0,1]]  # Обратите внимание, что в выводе программы # не должно быть повторяющихся троек.

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

1. Подход

Брутфорс-решение будет построено на итерациях. Представьте, что мы берем каждое число и сопоставляем его со всеми остальными числами, а затем сопоставляем каждую пару со всем другими числами. На что это похоже? Правильно, на вложенные циклы for. Вероятно, где-то здесь можно применить рекурсию. Если видите способ — поделитесь в комментариях!

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

2. Подготовка

Наш метод sum_of_three будет принимать два аргумента: список nums и число-сумму — target.

def sum_of_three(nums, target):

Далее нам нужно учесть крайний случай. Что, если в nums будет меньше 3 чисел? Тогда мы точно не сможем найти тройку чисел, дающих в сумме target. В этом случае мы вернем пустой список.

def sum_of_three(nums, target):     if len(nums) < 3:         return []

Теперь нам нужно инициализировать пару переменных, которые мы будем использовать. Во-первых, список output, в котором будут содержаться списки с найденными тройками чисел. Также мы создадим один из этих списков троек и назовем его current. Как только мы заполним current подходящей тройкой чисел, мы сможем добавить его в output.

output = [] current = []

И, как мы и говорили, нам нужно создать множество. Кроме того нам нужно будет сортировать списки чисел, чтобы не включить в множество один набор дважды. А такое может быть, если числа будут идти в разном порядке. Например, если у нас есть [1, -2, 1] и [1, 1, 2], нам нужно добавить во множество только какой-то один из этих списков.

checked = set() nums.sort()

3. Итерации

Итак, мы готовы к созданию циклов for. Всего нам нужно будет сделать три таких цикла. При этом самый важный момент — продумывание верхних и нижних границ для индексов.

В первом цикле мы начинаем с первого числа в списке и доходим до третьего с конца. Почему? Нам нужны два дополнительных числа, с которыми мы будем сопоставлять число в цикле, а когда мы дойдем до предпоследнего, у нас просто не останется достаточного количества чисел, чтобы составить из них тройку.

for i in range(len(nums) - 2):

В качестве напоминания: функция range() принимает два аргумента. Синтаксис:

range(start, stop, step)

Первые два аргумента задают диапазон, причем первый (start) включается в него, а второй (end) — не включается. Третий аргумент (step) опционален и задает шаг. Если он не указан, то по умолчанию шаг равен 1.

Но первый аргумент тоже опционален. Если его не указывать, по умолчанию стартовая позиция — 0.

Итак, что мы будем делать в нашем первом цикле for? План следующий:

1. Добавить первое число в current.
2. Вычесть это число из суммы.
3. Сверить с другими числами.
4. «Откатить» шаги 1 и 2, вернувшись в исходное положение, чтобы перейти к следующему числу.

Давайте реализуем это в коде. Первый шаг простой, мы лишь добавляем число под индексом i в current при помощи метода append.

current.append(nums[i])

Следующий шаг тоже довольно элементарный. Причем мы можем просто уменьшать target на текущее число (т. е. target -= nums[i]), но для лучшей читаемости давайте создадим новую переменную — working_target.

working_target = target - nums[i]

Третий шаг уже сложнее. Под «Сверить с другими числами» следует понимать вложенные циклы for, которыми мы займемся далее. Пока мы это просто закомментируем.

Четвертый шаг тоже простой. Мы удаляем последний элемент из current при помощи current.pop() и возвращаем target в исходный вид. Впрочем, здесь нам не придется ничего делать с target, потому что мы создали отдельную переменную. Наш цикл будет выглядеть так:

for i in range(len(nums) - 2):     current.append(nums[i])     working_target = target - nums[i]     # сделать вложенные циклы for     current.pop()

Отлично, теперь переходим ко второму, вложенному циклу.

Каков будет наш диапазон индексов? Мы начинаем с индекса, идущего следом за i, и доходим до предпоследней позиции в nums.

for j in range(i + 1, len(nums) - 1):

И повторяем шаги 1-4 из первого цикла. То есть мы добавляем число под индексом j в current, уменьшаем working_target на это число, оставляем место для третьего цикла, а затем откатываем назад шаги 1 и 2.

for j in range(i _ 1, len(nums) - 1):     working_target -= nums[j]     current.append(nums[j])     # сделать вложенный цикл for     current.pop()     working_target +=nums[j]

Наконец, давайте напишем самый внутренний цикл. С индексами уже должно быть понятно. Мы начинаем с j + 1 и идем до конца.

for k in range(j + 1, len(nums)):

Число под индексом k мы добавляем в current — тут никаких проблем нет. Но при вычитании k из working_target нам нужно проверить, равна ли разница нулю. Если все три числа в сумме дают target, то target — a — b — c = 0.

for k in range(j + 1, len(nums)):     current.append(nums[k])     if working_target - nums[k]  == 0:

Что у нас будет в блоке if? Пришла пора применить наше множество и проверить, нет ли в нем уже такой тройки чисел.

Мы не можем поместить списки во множество, так что нам нужно преобразовать нашу тройку чисел в строку. Следующий код создаст строку, в которой между числами будут проставлены знаки &. Например, [1, 1, -2] превратится в 1&1&-2. Для разделения чисел можно использовать любой символ, но он нужен обязательно — чтобы не спутать 1, 11 с 11, 1.

# добавление тройки чисел во множество для исключения дубликатов code = str(current[0]) for n in range(1, len(current)):     code += "&" + str(current[n])

Теперь мы можем проверить, есть ли code во множестве. Если его там нет, мы можем добавить его во множество и добавить current в наш итоговый output. Не забудьте, что добавить нужно копию current, а не просто ссылку. В противном случае тройка может измениться после добавления в output!

if code not in checked:     output.append(current.copy())     checked.add(code)

Наконец, мы удаляем третье число из current. В итоге третий цикл for выглядит следующим образом:

for k in range(j + 1, len(nums)):     current.append(nums[k])     if working_target - nums[k] == 0:         # добавление тройки чисел во множество для исключения дубликатов         code = str(current[0])         for n in range(1, len(current)):             code += "&" + str(current[n])         if code not in checked:             output.append(current.copy())             checked.add(code)     current.pop()

4. Заканчиваем работать над задачей «Сумма трех чисел»

Все, что нам остается, это вернуть список output. Наш метод в итоге выглядит так (следите внимательно за отступами во всех этих циклах for!):

def sum_of_three(nums, target):     if len(nums) < 3:         return []     output = []     current = []     checked = set()     nums.sort()      for i in range(len(nums) - 2):         current.append(nums[i])         working_target = target - nums[i]         for j in range(i + 1, len(nums) - 1):             working_target -= nums[j]             current.append(nums[j])             for k in range(j + 1, len(nums)):                 current.append(nums[k])                 if working_target - nums[k] == 0:                 # добавление тройки чисел во множество для исключения дубликатов                     code = str(current[0])                     for n in range(1, len(current)):                         code += "&" + str(current[n])                     if code not in checked:                         output.append(current.copy())                         checked.add(code)                 current.pop()             current.pop()             working_target += nums[j]         current.pop()         working_target += nums[i]     return output

Маленький пример для тестирования метода:

nums = [-1,0,1,2,-1,-4] target = 0 print(sum_of_three(nums, target))

Мы получим в итоговом output [[-1,-1,2],[-1,0,1]]. Обратите внимание, что у нас нет дубликатов троек, хотя во втором списке тоже есть две -1.

5. Наращиваем сложность: сумма четырех чисел

Теперь, когда вы счастливо перевели дух, покончив с суммой трех чисел, как насчет суммы четырех?

Хотя кажется, что эта задача сложнее, по сути, мы можем использовать все тот же брутфорс-метод. Я знаю, что временная сложность и так была ужасной — O(N3), а теперь станет еще хуже — O(N4). Но, как говорится, «зато работает».

Как мы подойдем к решению? Мы можем просто немного изменить код нашего метода sum_of_three, добавив еще один цикл. Каждый цикл начинается с предыдущего индекса, увеличенного на единицу, и заканчивается на единицу ближе к концу списка. Вот как это выглядит:

for i in range(len(nums) - 3):     for j in range(i + 1, len(nums) - 2):         for k in range(j + 1, len(nums) - 1):             for m in range(k + 1, len(nums)):

В каждом цикле мы проходим одинаковые шаги: добавляем число в current, вычитаем его из target, что-то проверяем, а затем откатываем назад предыдущие действия. Чтобы вам не пришлось лишний раз скроллить, вот вам наш готовый метод:

def sum_of_four(nums, target):     if len(nums) < 4:         return []     output = []     current = []     nums.sort()     checked = set()     for i in range(len(nums) - 3):         working_target = target - nums[i]         current.append(nums[i])         for j in range(i + 1, len(nums) - 2):             working_target -= nums[j]             current.append(nums[j])             for k in range(j + 1, len(nums) - 1):                 working_target -= nums[k]                 current.append(nums[k])                 for m in range(k + 1, len(nums)):                     current.append(nums[m])                     if working_target - nums[m] == 0:                         # добавление чисел во множество для исключения дубликатов                         code = str(current[0])                         for n in range(1, len(current)):                             code += "&" + str(current[n])                         if code not in checked:                             output.append(current.copy())                             checked.add(code)                     current.pop()                 current.pop()                 working_target += nums[k]             current.pop()             working_target += nums[j]         current.pop()     return output

Он точно такой же, как для «Суммы трех чисел», только с дополнительным циклом. Мы можем проверить его работу:

nums = [1,0,-1,0,-2,2, 0] target = 0 print(sum_of_four(nums, target))

В результате получим [[-2, -1, 1, 2], [-2, 0, 0, 2], [-1, 0, 0, 1]]. Работает!

6. Сумма пяти чисел и так далее

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

Да никаких проблем! Просто добавим еще один цикл for. Но сейчас вы можете заметить, что вырисовывается паттерн:

for j in range(i + 1, len(nums) - 2):     working_target -= nums[j]     current.append(nums[j])     for k in range(j + 1, len(nums) - 1):         working_target -= nums[k]         current.append(nums[k])         for m in range(k + 1, len(nums)):             current.append(nums[m])

А когда у нас есть повторяющийся код, хороший тон — написать для него отдельный метод. Почему бы нам не написать рекурсивный метод для запуска каждого цикла for? Тогда мы сможем просто указывать ему, сколько чисел мы хотим суммировать, а он сам определит, сколько вложенных циклов сделать.

Итак, давайте сделаем такой же метод, как раньше, только теперь он будет принимать еще и аргумент k — количество чисел, сумма которых должна быть равна target (три, четыре, пять и т. д. чисел).

def sum_of_k(nums, target, k):     if len(nums) < k:         return []     output = []     current = []     nums.sort()     checked = set()

А теперь мы создадим рекурсивный метод sum_helper(), который будет принимать… да практически все переменные, созданные нам ранее. Кроме того, нам нужно знать предыдущий индекс, стартующий с 0 (чтобы начинать цикл с позиции i + 1). Этот метод будет изменять массивы, поэтому нам нужно вызвать его, а затем вернуть output.

sum_helper(nums, target, k, output, checked, current, 0) return output

7. Итерация с рекурсивным методом

Давайте подумаем, как нам создать наш рекурсивный вспомогательный метод. Мы каждый раз будем вычитать из k. В конечном итоге, «сумма четырех» — это просто «сумма трех» с дополнительным числом.

Нашему рекурсивному методы нужен базовый случай. Если k равна 0, нам не нужно ничего делать. В этом случае мы просто выходим из функции (return).

if k == 0:     return

Если k > 0, нам нужно перебрать в цикле nums. Мы начинаем с индекса prev и идем до «конец списка — k + 1» (к этому мы пришли, рассматривая паттерн в предыдущем решении).

for i in range(prev, len(nums) - k + 1):

Что дальше? Смотрим на наш список: мы добавляем текущее число в current.

current.append(nums[i])

Возможно, вы думаете, что следующий шаг — вычесть это число из target, но сначала нам нужно кое-что проверить. Этот цикл for должен быть универсальным для каждой итерации k, так что мы должны учесть и случай, если данный цикл — самый вложенный.

То есть нам нужно проверить, дает ли нуль вычитание числа из target. Затем мы делаем то же, что и раньше (можете скопировать с минимальными правками) — добавляем code в виде строки во множество и, если значение уникально, добавляем копию current в output.

if k == 1 and target - nums[i] == 0:     # добавление чисел во множество для исключения дубликатов     code = str(current[0])     for n in range(1, len(current)):         code += "&" + str(current[n])     if code not in checked:         output.append(current.copy())         checked.add(code)

Теперь нам нужно вычесть число из рабочей target и написать наш рекурсивный случай. Но погодите, мы же можем сделать это за один шаг! Вот рекурсивный вызов:

sum_helper(nums, target - nums[i], k - 1, output, checked, current, i + 1)

Здесь у нас меняется target, k уменьшается на 1, а текущий индекс + 1 становится предыдущим индексом prev.

Наконец, нам нужно откатить назад добавление числа в current:

current.pop()

Вот и все! Вместе наши методы выглядят следующим образом:

def sum_of_k(nums, target, k):     if len(nums) < k:         return []     output = []     current = []     nums.sort()     checked = set()     sum_helper(nums, target, k, output, checked, current, 0)     return output  def sum_helper(nums, target, k, output, checked, current, prev):     if k == 0:         return     for i in range(prev, len(nums) - k + 1):         current.append(nums[i])         if k == 1 and target - nums[i] == 0:             # добавление чисел во множество для исключения дубликатов             code = str(current[0])             for n in range(1, len(current)):                 code += "&" + str(current[n])             if code not in checked:                 output.append(current.copy())                 checked.add(code)         sum_helper(nums, target - nums[i], k - 1, output, checked, current, i + 1)         current.pop()

Наш метод по-прежнему очень громоздкий, но нам хотя бы не приходится изменять его вручную при каждом изменении k. Следующий код будет работать независимо от того, равна k 3, 4 или другим числам!

nums = [-1,0,1,2,-1,-4, -2] target = 0 k = 4 print(sum_of_k(nums, target, k)) # -> [[-2, -1, 1, 2], [-1, -1, 0, 2]]

Не стоит забывать, что это — брутфорс-решение. Его временная сложность — O(Nk), а это очень плохо, если k оказывается большим числом. Вероятно, эту задачу можно решить более элегантно. Если вы знаете, как, — добро пожаловать в комментарии!