三数之和的多种可能

标签: 数组 哈希表 双指针 计数 排序

难度: Medium

给定一个整数数组 arr ,以及一个整数 target 作为目标值,返回满足 i < j < k arr[i] + arr[j] + arr[k] == target 的元组 i, j, k 的数量。

由于结果会非常大,请返回 109 + 7 的模。

示例 1:

输入:arr = [1,1,2,2,3,3,4,4,5,5], target = 8
输出:20
解释:
按值枚举(arr[i], arr[j], arr[k]):
(1, 2, 5) 出现 8 次;
(1, 3, 4) 出现 8 次;
(2, 2, 4) 出现 2 次;
(2, 3, 3) 出现 2 次。

示例 2:

输入:arr = [1,1,2,2,2,2], target = 5
输出:12
解释:
arr[i] = 1, arr[j] = arr[k] = 2 出现 12 次:
我们从 [1,1] 中选择一个 1,有 2 种情况,
从 [2,2,2,2] 中选出两个 2,有 6 种情况。

提示:

  • 3 <= arr.length <= 3000
  • 0 <= arr[i] <= 100
  • 0 <= target <= 300

Submission

运行时间: 38 ms

内存: 16.1 MB

class Solution:
    def threeSumMulti(self, arr: List[int], target: int) -> int:
        modi = 10 ** 9 + 7
        cnt = Counter(arr)
        nums = sorted(cnt.keys())
        ans= 0

        for i in nums:
            for j in nums:
                k = target-i-j
                if i > j or j > k: continue  # 仅一条,排除错误选项
                if i == j == k:
                    ans += comb(cnt[i], 3)  # 不对cnt[i]是否 >= 3 进行判断
                elif i == j:
                    ans += cnt[k] * comb(cnt[i], 2)  # 不对 k 是否在 cnt 中进行特判
                elif j == k:
                    ans += cnt[i] * comb(cnt[j], 2)
                else:
                    ans += cnt[k] * cnt[i] * cnt[j]
                
        return ans%modi

Explain

此题解采用哈希表和组合计数的方法来解决三数之和的问题。首先,使用Counter统计数组中每个数字的出现次数。然后,对哈希表的键进行排序,以便在后续的双重循环中按序处理。在双重循环中,对于每一对数字(i, j),计算第三个数字k为target-i-j。如果i,j和k满足条件i <= j <= k,则根据i,j和k的关系(是否相等)使用组合公式计算满足条件的三元组数量,并对结果进行累加。特别注意,只有当k也在哈希表中时,才会添加到答案中。最终,返回累加结果对1e9+7取模的值。

时间复杂度: O(1)

空间复杂度: O(1)

from collections import Counter
from math import comb

class Solution:
    def threeSumMulti(self, arr: List[int], target: int) -> int:
        modi = 10 ** 9 + 7
        cnt = Counter(arr)  # 统计每个数出现的次数
        nums = sorted(cnt.keys())  # 对数组中出现的数字进行排序
        ans = 0

        for i in nums:
            for j in nums:
                k = target - i - j  # 计算第三个数
                if i > j or j > k: continue  # 确保i <= j <= k
                if i == j == k:
                    ans += comb(cnt[i], 3)  # 当三个数相同时,使用组合数计算
                elif i == j:
                    ans += cnt[k] * comb(cnt[i], 2)  # 当两个数相同时,另一个数为k
                elif j == k:
                    ans += cnt[i] * comb(cnt[j], 2)  # 当两个数相同时,另一个数为i
                else:
                    ans += cnt[k] * cnt[i] * cnt[j]  # 三个数均不相同

        return ans % modi  # 返回结果对10^9+7取模

Explore

组合公式 comb(n, k) 计算的是从 n 个元素中选取 k 个元素的组合数。当 i, j, k 相等时,意味着我们需要从数 i 的出现次数(cnt[i])中选取3个。例如,如果 i 出现了5次,那么从这5次中挑选3个可以形成的组合数为 comb(5, 3)。这是因为每个组合代表一种不同的方式将这些相同的数分配到三个位置上(不考虑顺序)。使用 comb(cnt[i], 3) 是为了确保这种所有可能的选择都被计算在内。

在双重循环中,我们保持条件 i <= j <= k 来避免重复和非必要的计算。这种筛选确保每个三元组仅被考虑一次,防止例如 (1,2,3)、(2,1,3)、(3,1,2) 这样相同的三元组以不同顺序出现多次。此外,这也简化了逻辑,因为我们不需要在后续的代码中检查和管理这些重复的情况。如果 k < j,那么就会与已经设定的 i <= j 相矛盾,因此这样的组合可以直接跳过。

如果在计算出 k = target - i - j 后发现 k 不在哈希表中,这意味着在数组中不存在这样的数来与 i 和 j 配对以形成符合条件的三元组。因此,这种情况下,我们应直接跳过该组合,并继续寻找其他可能的 i 和 j 组合。此处理是必须的,因为只有当三个数都存在时,才能形成有效的三元组。

对 10^9 + 7 取模是常用的技术,用于防止在进行大量计算时结果超出常规整数范围而导致溢出。10^9 + 7 是一个大的质数,也是常用的模数,因为它能确保运算结果保持在一个安全的范围内,同时这个数也足够大,使得在模运算下结果还能保持良好的随机性和均匀分布。此外,在编程竞赛和算法实现中,这种模运算还可以加快计算速度和增强代码的鲁棒性。