三数之和问题 | 豆包MarsCode AI刷题

问题描述

小U有一个整数数组 arr，他希望找到其中三个元素 i, j, k 满足条件 i < j < k 且 arr[i] + arr[j] + arr[k] == target。由于可能存在大量的元组，结果需要对 10^9 + 7 取模。

例如：当 arr = [1,1,2,2,3,3,4,4,5,5] 且 t = 8 时，有多个元组满足条件，你需要返回这些元组的数量。

题目分析

我们需要在一个整数数组 arr 中找到满足以下条件的三元组 (i, j, k)：

1. 下标关系：i < j < k。
2. 数值关系：arr[i] + arr[j] + arr[k] == target。

由于数组可能很大，暴力枚举 O(n^3) 的复杂度不可行，因此我们需要优化解法。同时，结果可能很大，需要对 (10^9 + 7) 取模。

---

解题思路

为了解决这个问题，可以采用双指针法，结合排序数组和巧妙的计数方式，降低时间复杂度。

1. 排序数组

首先对数组 arr 进行排序，这样可以根据当前选取的元素之和与目标数的大小关系，选择减小元素或增大元素，方便地使用双指针法在数组的子区间中查找满足条件的元素。

2. 固定第一个元素

通过一个循环，逐一固定每个元素为当前三元组的第一个元素 arr[i]。对于每个 i：

• 剩余的两个元素 arr[j] 和 arr[k] 的和需要满足 arr[j] + arr[k] == target - arr[i]。

3. 双指针查找

在固定 arr[i] 的情况下，使用双指针在 i+1 到arr.length - 1查找两个数：

• 初始化左指针 j = i + 1 和右指针 k = arr.length - 1。
• 计算三数之和 sum = arr[i] + arr[j] + arr[k]：
  • 如果 sum > target，说明右指针太大，k--。
  • 如果 sum < target，说明左指针太小，j++。
  • 如果 sum == target：
    • 处理重复值：
      • 如果 arr[j] == arr[k]，说明从 j 到 k 的所有值都相同，可以直接计算组合数：
        $$\text{count} = \frac{(k - j + 1) \times (k - j)}{2}$$
        并将其加入结果。
      • 否则，分别统计 arr[j] 和 arr[k] 的重复次数 left 和 right，结果为 left × right。

4. 结果取模

为避免结果溢出，所有的结果累加都需要取模 (10^9 + 7)。

---

代码实现

以下是完整的代码：

import java.util.*;
public class Main {
    public static int solution(int[] arr, int t) {
        Arrays.sort(arr);  // 排序数组
        int ans = 0;
        int mod = 1000000007;
        for (int i = 0; i < arr.length - 2; i++) {
            int j = i + 1, k = arr.length - 1;  // 初始化双指针
            while (j < k) {
                int sum = arr[i] + arr[j] + arr[k];
                if (sum > t) {
                    k--;  // 如果和太大，右指针左移
                } else if (sum < t) {
                    j++;  // 如果和太小，左指针右移
                } else {
                    if (arr[j] == arr[k]) {
                        // 处理 arr[j] 和 arr[k] 相等的情况
                        int count = k - j + 1;
                        ans = (ans + count * (count - 1) / 2) % mod;
                        break;
                    } else {
                        // 分别统计 arr[j] 和 arr[k] 的重复次数
                        int left = 1, right = 1;
                        while (j + left < k && arr[j + left] == arr[j]) {
                            left++;
                        }
                        while (k - right > j && arr[k - right] == arr[k]) {
                            right++;
                        }
                        // 累加结果
                        ans = (ans + left * right) % mod;
                        j += left;
                        k -= right;
                    }
                }
            }
        }
        return ans;
    }

    public static void main(String[] args) {
        System.out.println(solution(new int[]{1, 1, 2, 2, 3, 3, 4, 4, 5, 5}, 8) == 20);
        System.out.println(solution(new int[]{2, 2, 2, 2}, 6) == 4);
        System.out.println(solution(new int[]{1, 2, 3, 4, 5}, 9) == 2);
        System.out.println(solution(new int[]{1, 1, 1, 1}, 3) == 4);
    }
}

---

示例分析

示例 1

输入：arr = [1, 1, 2, 2, 3, 3, 4, 4, 5, 5]，target = 8

1. 排序后数组为 [1, 1, 2, 2, 3, 3, 4, 4, 5, 5]。
2. 通过固定每个 arr[i]，使用双指针法查找 (j, k)：
   • 固定 arr[i] = 1，找到 (1, 2, 5) 和 (1, 3, 4) 的多种组合。
   • 固定 arr[i] = 2，找到 (2, 2, 4) 的多种组合。
3. 总共有 20 种组合。

输出：20。

示例 2

输入：arr = [2, 2, 2, 2]，target = 6

1. 排序后数组为 [2, 2, 2, 2]。
2. 通过固定 arr[i] = 2，找到 (2, 2, 2) 的 4 种组合。

输出：4。

---

总结

• 本题通过排序和双指针法有效降低了时间复杂度。
• 对重复值的处理是关键，利用排序后的数组巧妙统计组合数，大大提高了效率。
• 代码清晰简洁，适合处理类似三数和问题，如：(15. 三数之和 - 力扣（LeetCode）)