使用itertools的Python 3Sum解法超时，如何优化？

原解法的问题分析 #

你的解法用itertools.combinations生成所有可能的三元组，本质是暴力枚举，时间复杂度达到O(n³)——当数组元素数量较大时（比如n=1000），生成的三元组数量会超过1.6亿，必然超时。另外，事后用i not in z去重，每次判断都要遍历列表，进一步增加了时间成本。

优化思路：排序+双指针 #

这是解决三数之和问题的经典高效解法，时间复杂度降到O(n²)，且能在搜索过程中直接避免重复三元组：

1. 先排序数组：排序后可以利用单调性缩小搜索范围，还能轻松跳过重复元素。
2. 固定第一个元素，双指针找另外两个数：
   • 遍历每个元素作为三元组的第一个数，若当前元素和前一个元素相同，直接跳过（避免生成重复三元组）。
   • 左指针从当前元素的下一位开始，右指针从数组末尾开始，根据三者的和调整指针位置：
     • 和为0：记录三元组，然后同时移动左右指针，并且跳过所有重复的元素，避免重复记录。
     • 和小于0：左指针右移，增大总和。
     • 和大于0：右指针左移，减小总和。

优化后的代码 #

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        result = []
        n = len(nums)
        for i in range(n):
            # 跳过重复的第一个元素
            if i > 0 and nums[i] == nums[i-1]:
                continue
            left = i + 1
            right = n - 1
            while left < right:
                total = nums[i] + nums[left] + nums[right]
                if total == 0:
                    result.append([nums[i], nums[left], nums[right]])
                    # 跳过左指针重复元素
                    while left < right and nums[left] == nums[left+1]:
                        left += 1
                    # 跳过右指针重复元素
                    while left < right and nums[right] == nums[right-1]:
                        right -= 1
                    left += 1
                    right -= 1
                elif total < 0:
                    left += 1
                else:
                    right -= 1
        return result

为什么这个解法更高效 #

• 排序的时间是O(n log n)，双指针遍历的时间是O(n²)，整体时间复杂度远低于暴力枚举，能轻松通过所有测试用例。
• 在搜索过程中直接跳过重复元素，不需要事后再做去重操作，节省了额外的时间开销。

内容的提问来源于stack exchange，提问作者Mubarak Ajibola Odufade