Задача: 923. 3Sum With Multiplicity
Сложность: medium

Если задан целочисленный массив arr и целое число target, верните количество кортежей i, j, k, таких, что i < j < k и arr[i] + arr[j] + arr[k] == target. Поскольку ответ может быть очень большим, верните его по модулю 10^9 + 7.

Пример:

Input: arr = [1,1,2,2,3,3,4,4,5,5], target = 8
Output: 20

👨‍💻 Алгоритм:

1⃣Отсортировать массив arr.

2⃣Инициализировать счетчик для количества кортежей.
Пройти по массиву тремя указателями i, j, и k:
Для каждого i, установить j на i + 1, и k на конец массива.
Использовать двухуказательный метод для нахождения пар (j, k), таких что arr[i] + arr[j] + arr[k] == target.

3⃣Вернуть результат по модулю 10^9 + 7.

😎 Решение:

class Solution {
    func threeSumMulti(_ arr: [Int], _ target: Int) -> Int {
        let MOD = 1_000_000_007
        var count = 0
        let arr = arr.sorted()
        
        for i in 0..<arr.count {
            var j = i + 1
            var k = arr.count - 1
            
            while j < k {
                let sum = arr[i] + arr[j] + arr[k]
                if sum == target {
                    if arr[j] == arr[k] {
                        count += (k - j + 1) * (k - j) / 2
                        break
                    } else {
                        var left = 1
                        var right = 1
                        while j + 1 < k && arr[j] == arr[j + 1] {
                            left += 1
                            j += 1
                        }
                        while k - 1 > j && arr[k] == arr[k - 1] {
                            right += 1
                            k -= 1
                        }
                        count += left * right
                        j += 1
                        k -= 1
                    }
                } else if sum < target {
                    j += 1
                } else {
                    k -= 1
                }
            }
        }
        
        return count % MOD
    }
}

Ставь 👍 и забирай 📚 Базу знаний