322三数之和问题

题解

该问题要求在一个整数数组中找到所有满足 i < j < k 且 arr[i] + arr[j] + arr[k] == target 的三元组，并返回这些三元组的数量。由于可能存在大量的组合，结果需要对 $10^9 + 7$ 取模。

解题思路

1. 排序数组：首先将数组进行排序，方便使用双指针法。
2. 遍历数组：固定第一个数 arr[i]，然后使用双指针 j 和 k 来查找剩下的两个数。
3. 双指针查找：
   • 如果 arr[i] + arr[j] + arr[k] 小于目标值，移动左指针 j 以增大和。
   • 如果大于目标值，移动右指针 k 以减小和。
   • 当找到等于目标值的三元组时，需处理可能存在的重复元素：
     • 如果 arr[j] 和 arr[k] 相同，则可以从这段区间中任意选择两个数字组合。可以组合出 (k - j + 1) * (k - j) / 2 个三元组。
     • 否则，计算左右指针处相同元素的个数 l 和 r，并将 l * r 加入结果中。
4. 取模：在每次更新结果时，对结果取模以防止溢出。

代码

mod = int(1e9 + 7)
def solution(arr: list, t: int) -> int:
    arr.sort()
    n = len(arr)
    res = 0
    for i in range(n - 2):
        j, k = i + 1, n - 1
        while j < k:
            if arr[i] + arr[j] + arr[k] < t:
                j += 1
            elif arr[i] + arr[j] + arr[k] > t:
                k -= 1
            else:
                if arr[j] == arr[k]:
                    res += (k - j + 1) * (k - j) // 2
                    break
                l = r = 1
                while j < k and arr[j] == arr[j + 1]:
                    l += 1
                    j += 1
                while j < k and arr[k] == arr[k - 1]:
                    k -= 1
                    r += 1
                j += 1
                k -= 1
                res = (res + l * r) % mod
    return res

if __name__ == '__main__':
    print(solution([1, 1, 2, 2, 3, 3, 4, 4, 5, 5], 8) == 20)
    print(solution([2, 2, 2, 2], 6) == 4)
    print(solution([1, 2, 3, 4, 5], 9) == 2)
    print(solution([1, 1, 1, 1], 3) == 4)

复杂度分析

• 时间复杂度：$O(n^2)$，其中 $n$ 是数组的长度。排序需要 $O(n \log n)$，双指针遍历需要 $O(n^2)$。
• 空间复杂度：$O(1)$，只使用了常数级别的额外空间。