Алгоритминутка

Дан список целых чисел nums. Вернуть все тройки чисел [nums[i], nums[j], nums[k]], такие что i != j,i != k, j != k и nums[i] + nums[j] + nums[k] = 0. Имеется ввиду неповторяющиеся тройки, то есть тройки [-1, 0, 1] и [0, -1, 1] считаем одинаковыми.

Примеры:

1) nums = [-1,0,1,2,-1,-4] —> [[-1,-1,2],[-1,0,1]]
2) nums = [0,1,1] —> []
3) nums = [0,0,0] —> [[0,0,0]]

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

Решение №1

0) Заводим список answer, в котором будем хранить подходящие тройки чисел. Запустим три цикла по списку с индексами i, j, k

1) Если мы нашли тройку nums[i] + nums[j] + nums[k] = 0 и все индексы различны, то сортируем ее и проверяем есть ли такая в answer. Если есть, то пропускаем, если нет, то добавляем. Сортировать нужно, чтобы в ответе не было одинаковых троек.

2) После завершения циклов возвращаем answer.

def solve(nums):
  answer = []

  for i in range(len(nums)):
    for j in range(i+1, len(nums)):
      for k in range(j+1, len(nums)):
        if nums[i] + nums[j] + nums[k] == 0:
          candidate = sorted([nums[i], nums[j], nums[k]])
          if candidate not in answer:
            answer.append(candidate)
  return answer

Решение №2

0) Заводим список answer, в котором будем хранить подходящие тройки чисел. Сортируем список nums.

1) Запускаем цикл по списку с индексом i. Вводим два дополнительных указателя: l = i+1 и r = n-1 (n — длина списка nums). 

2) Далее перебираем тройки nums[i] + nums[l] + nums[r] пока l будет меньше r. Если это сумма меньше 0, то двигаем l на 1 вправо, если больше нуля, то двигаем r на 1 влево. 

3) Если же эта сумма равна нулю, то сохраняем эту тройку в answer. Но нужно рассмотреть случай, если nums[l] = nums[l+1], тогда на следующей итерации добавиться такая же тройка. Чтобы этого избежать будем отдельным циклом сдвигать индекс l пока не встретим отличный от этого элемент. Аналогично и с индексом r.

4) Так как список отсортированный, то после того как мы встретили положительное число nums[i] > 0, цикл можно обрывать (тройки с отрицательными числами уже проверили)

def solve(nums):
  nums.sort()
  answer = []
  n = len(nums)

  for i in range(n - 2):
    if nums[i] > 0:
      break
    if i > 0 and nums[i] == nums[i-1]:
      continue
    l = i + 1
    r = n - 1
    while l < r:
      total = nums[i] + nums[l] + nums[r]
      if total < 0:
        l += 1
      elif total > 0:
        r -= 1
      else:
        triplet = [nums[i], nums[l], nums[r]]
        answer.append(triplet)
        while l < r and nums[l] == triplet[1]:
          l += 1
        while l < r and nums[r] == triplet[2]:
          r -= 1
  return answer