Как найти сумму трех чисел в питоне

Давайте рассмотрим довольно классическую задачку на программирование под названием «Сумма трех чисел» (и производную от нее — «Сумму четырех чисел»). Решать будем брутфорс-методом, а затем усовершенствуем решение при помощи рекурсии так, чтобы оно подходило для любой задачи «Сумма 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(N³), а теперь станет еще хуже — O(N⁴). Но, как говорится, «зато работает».

Как мы подойдем к решению? Мы можем просто немного изменить код нашего метода 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(N^k), а это очень плохо, если k оказывается большим числом. Вероятно, эту задачу можно решить более элегантно. Если вы знаете, как, — добро пожаловать в комментарии!