Триплеты с суммой меньше целевой
Условие задачи
Даны неотсортированный массив arr и целевая сумма. Подсчитайте все триплеты разных индексов i, j и k, для которых arr[i] + arr[j] + arr[k] < target.
Пример 1:
1
2
3
Вход: [-1, 0, 2, 3], target=3
Выход: 2
Пояснение: подходят [-1, 0, 3] и [-1, 0, 2].
Пример 2:
1
2
3
Вход: [-1, 4, 2, 1, 3], target=5
Выход: 4
Пояснение: подходят [-1, 1, 4], [-1, 1, 3], [-1, 1, 2] и [-1, 2, 3].
Решение
Отсортируем массив и будем фиксировать первый элемент триплета. Для оставшейся пары используем указатели на начало и конец ещё не просмотренной части массива.
Если arr[left] + arr[right] меньше остатка целевой суммы, то подходят все пары текущего left с элементами от left+1 до right: их right-left. После этого увеличиваем left. Если сумма слишком велика, уменьшаем right.
Код
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
package main
import (
"fmt"
"sort"
)
func countTripletsWithSmallerSum(arr []int, target int) int {
sort.Ints(arr)
count := 0
for i := 0; i < len(arr)-2; i++ {
left, right := i+1, len(arr)-1
pairTarget := target - arr[i]
for left < right {
if arr[left]+arr[right] < pairTarget {
count += right - left
left++
} else {
right--
}
}
}
return count
}
func main() {
fmt.Println(countTripletsWithSmallerSum([]int{-1, 0, 2, 3}, 3))
fmt.Println(countTripletsWithSmallerSum([]int{-1, 4, 2, 1, 3}, 5))
}
Вывод:
1
2
2
4
Временная сложность
Сортировка занимает \(O(N \log N)\). Для каждого первого элемента два указателя проходят оставшуюся часть массива, поэтому общая временная сложность равна \(O(N^2)\).
Пространственная сложность
Если не учитывать результат, сортировке может потребоваться до \(O(N)\) дополнительной памяти.
Вариация задачи
Если вместо количества нужно вернуть сами триплеты, при найденной меньшей сумме добавляем пары текущего left со всеми элементами от right до left+1. В худшем случае вывод содержит \(O(N^3)\) триплетов, поэтому время и память результата также достигают \(O(N^3)\).