Давайте рассмотрим довольно классическую задачку на программирование под названием «Сумма трех чисел» (и производную от нее — «Сумму четырех чисел»). Решать будем брутфорс-методом, а затем усовершенствуем решение при помощи рекурсии так, чтобы оно подходило для любой задачи «Сумма 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? План следующий:
- Добавить первое число в
current. - Вычесть это число из суммы.
- Сверить с другими числами.
- «Откатить» шаги 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 оказывается большим числом. Вероятно, эту задачу можно решить более элегантно. Если вы знаете, как, — добро пожаловать в комментарии!
