Триплеты с нулевой суммой

Триплеты с нулевой суммой


Условие задачи

Дан неотсортированный массив чисел. Найдите в нём все уникальные триплеты, сумма которых равна нулю.

Пример 1:

1
2
3
Вход: [-3, 0, 1, 2, -1, 1, -2]
Выход: [-3, 1, 2], [-2, 0, 2], [-2, 1, 1], [-1, 0, 1]
Пояснение: существуют четыре уникальных триплета с нулевой суммой.

Пример 2:

1
2
3
Вход: [-5, 2, -1, -2, 3]
Выход: [-5, 2, 3], [-2, -1, 3]
Пояснение: существуют два уникальных триплета с нулевой суммой.

Решение

Сначала отсортируем массив. Затем будем по очереди фиксировать число X и искать справа от него пару Y и Z, для которой Y + Z = -X.

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

Код

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
package main

import (
	"fmt"
	"sort"
)

func searchTriplets(arr []int) [][]int {
	sort.Ints(arr)
	triplets := [][]int{}

	for i := 0; i < len(arr)-2; i++ {
		if i > 0 && arr[i] == arr[i-1] {
			continue
		}

		left, right := i+1, len(arr)-1
		for left < right {
			currentSum := arr[left] + arr[right]
			targetSum := -arr[i]
			switch {
			case currentSum == targetSum:
				triplets = append(triplets, []int{arr[i], arr[left], arr[right]})
				left++
				right--
				for left < right && arr[left] == arr[left-1] {
					left++
				}
				for left < right && arr[right] == arr[right+1] {
					right--
				}
			case currentSum < targetSum:
				left++
			default:
				right--
			}
		}
	}

	return triplets
}

func main() {
	fmt.Println(searchTriplets([]int{-3, 0, 1, 2, -1, 1, -2}))
	fmt.Println(searchTriplets([]int{-5, 2, -1, -2, 3}))
}

Вывод:

1
2
[[-3 1 2] [-2 0 2] [-2 1 1] [-1 0 1]]
[[-5 2 3] [-2 -1 3]]

Временная сложность

Сортировка занимает \(O(N \log N)\), а поиск пары запускается для каждого элемента и занимает \(O(N)\). Общая временная сложность равна \(O(N \log N + N^2)\), то есть \(O(N^2)\).

Пространственная сложность

Если не учитывать результат, сортировке может потребоваться до \(O(N)\) дополнительной памяти. Размер результата зависит от количества найденных уникальных триплетов.