Python中如何优化LeetCode四数之和算法题

https://leetcode-cn.com/problems/4sum/description/

我的代码，提交解答的时候通不过，显示超出时间限制，请大佬帮我优化一下：

class Solution:
    def fourSum(self, nums, target):
        """
        :type nums: List[int]
        :type target: int
        :rtype: List[List[int]]
        """
        res, dicti = set(), {}
        numLen = len(nums)
        nums.sort()
        for i in range(numLen):
            for j in range(i+1, numLen):
                key = nums[i] + nums[j]
                if key not in dicti.keys():
                    dicti[key] = [(i,j)]
                else:
                    dicti[key].append((i,j))
        for i in range(numLen):
            for j in range(i+1, numLen-2):
                exp = target - nums[i] -nums[j]
                if exp in dicti.keys():
                    for tmpIndex in dicti[exp]:
                        if tmpIndex[0] > j:
                            res.add((nums[i], nums[j], nums[tmpIndex[0]], nums[tmpIndex[1]]))
        return [list(i) for i in res]

Python中如何优化LeetCode四数之和算法题

4 回复

>>> nums = [1, 0, -1, 0, -2, 2]
>>> target = 0
>>> from itertools import combinations as cb
>>> [c for c in cb(nums, 4) if sum©==target]
[(1, 0, -1, 0), (1, -1, -2, 2), (0, 0, -2, 2)]
>>>


对于LeetCode四数之和问题，核心优化思路是在三数之和基础上增加一层循环，同时利用排序和双指针技巧。最直接的暴力解法是O(n^4)，优化后能达到O(n^3)。

这里给出一个经过剪枝优化的标准解法：

def fourSum(nums, target):
    nums.sort()
    n = len(nums)
    res = []
    
    for i in range(n - 3):
        # 一级剪枝
        if nums[i] * 4 > target:
            break
        if i > 0 and nums[i] == nums[i - 1]:
            continue
            
        for j in range(i + 1, n - 2):
            # 二级剪枝
            if nums[i] + nums[j] * 3 > target:
                break
            if j > i + 1 and nums[j] == nums[j - 1]:
                continue
                
            left, right = j + 1, n - 1
            while left < right:
                total = nums[i] + nums[j] + nums[left] + nums[right]
                
                if total == target:
                    res.append([nums[i], nums[j], 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 < target:
                    left += 1
                else:
                    right -= 1
                    
    return res

关键优化点：

  1. 先排序数组，这样可以用双指针将内层两数之和从O(n^2)降到O(n)
  2. 两层循环中都有提前剪枝：当当前最小值组合已超过target时直接break
  3. 跳过重复元素避免重复结果

这样整体时间复杂度O(n^3)，空间复杂度O(1)（不考虑结果存储）。对于n=200的数据量完全够用。

总结：排序+双指针+剪枝是标准解法。

dicti.keys()返回的是一个 list，判断一个元素是不是在一个 list 里是 O(n)的。这里直接 if exp in dicti 就可以了，判断一个元素是不是在一个 dict 里是 O(1)的。

建议 lz 搞清楚基本数据结构再写

回到顶部