Python高效解决LeetCode三数之和问题：从超时到O(N^2)优化实践

本文深入探讨了leetcode三数之和（3sum）问题的高效python解法。针对常见的超时问题，文章将详细分析原始解法的性能瓶颈，并介绍如何通过数组排序与双指针技术，将时间复杂度从低效优化至o(n^2)。教程涵盖了算法原理、代码实现以及关键的去重策略，旨在帮助读者掌握解决此类问题的最佳实践。

理解三数之和问题

三数之和（3Sum）问题要求从一个整数数组 nums 中找出所有唯一的三元组 [nums[i], nums[j], nums[k]]，使得 i != j、i != k、j != k，并且 nums[i] + nums[j] + nums[k] == 0。需要注意的是，最终的解集中不能包含重复的三元组。

这个问题在算法面试中非常常见，它考验了开发者对数组操作、去重逻辑以及时间复杂度的优化能力。

原始解法的性能瓶颈分析

许多初学者在解决三数之和问题时，可能会采用一种直观但效率不高的策略，导致“时间超出限制”（Time Limit Exceeded）。以下是常见的一种低效解法示例：

def threeSum(nums):
    sol = []
    pos = 1
    nums.sort()

    def search(p, vals):
        l, r = 0, len(vals) - 1
        sols = []
        while l < p < r:
            current_sum = vals[l] + vals[p] + vals[r]
            if current_sum == 0:
                sols.append([vals[l], vals[p], vals[r]])
                # 列表pop操作的代价很高
                vals.pop(r) 
                vals.pop(l)
                l, r = l, r - 2
                p -= 1
                continue
            if current_sum > 0:
                r -= 1
            if current_sum < 0:
                l += 1
        return sols

    while pos < len(nums) - 1:
        # 每次都创建新的列表切片，且列表in操作代价高
        new_sol = search(pos, nums[:]) 
        for n in new_sol:
            if n not in sol: # 检查重复三元组的代价高
                sol.append(n)
        pos += 1
    return sol

该解法首先对数组进行了排序，这是一个良好的开端。然而，其核心的 search 函数以及主循环中存在几个严重的性能问题：

立即学习“Python免费学习笔记（深入）”；

1. list.pop() 操作: 在 search 函数内部，当找到一个三元组后，会使用 vals.pop(r) 和 vals.pop(l) 从列表中移除元素。Python 列表的 pop() 操作，特别是当移除的不是末尾元素时（例如 pop(l)），需要移动后续所有元素，其时间复杂度为 O(N)，其中 N 是当前列表的长度。在循环中多次执行此类操作，会显著增加整体运行时间。
2. 列表切片 nums[:]: 在主循环中，每次调用 search 函数时都创建了 nums[:] 的副本。这会产生 O(N) 的空间和时间开销，并且在每次迭代中重复进行。
3. n not in sol 检查: 为了避免重复的三元组，代码使用了 if n not in sol: 进行检查。当 sol 是一个列表时，in 操作的时间复杂度为 O(M)，其中 M 是 sol 中元素的数量。如果 sol 包含大量三元组，每次检查都会耗费大量时间，可能导致总复杂度接近 O(N^4)。

综合来看，这种方法的时间复杂度远高于 O(N^2)，尤其是在处理大型测试用例时，很容易超时。

ViiTor实时翻译

AI实时多语言翻译专家！强大的语音识别、AR翻译功能。

116

查看详情

高效的O(N^2)解法：排序与双指针

解决三数之和问题的标准高效方法是结合排序和双指针技术。这种方法能够将时间复杂度优化到 O(N^2)。

核心思想

1. 排序: 首先对输入数组 nums 进行排序。排序是此方法的基础，它使得我们可以利用元素的有序性来高效地查找和跳过重复项。排序的时间复杂度为 O(N log N)。
2. 固定一个元素: 遍历排序后的数组，用一个指针 i 固定第一个元素 nums[i]。
3. 双指针查找: 对于每一个固定的 nums[i]，我们实际上是将问题转化为了“在 nums[i+1:] 子数组中找到两个数 nums[lo] 和 nums[hi]，使得 nums[lo] + nums[hi] == -nums[i]”。这正是经典的“两数之和”问题，可以在一个已排序的数组中使用双指针技术高效解决。
   • 设置左指针 lo 从 i+1 开始，右指针 hi 从数组末尾 len(nums) - 1 开始。
   • 在 lo < hi 的条件下循环：
     • 计算当前三数之和 current_sum = nums[i] + nums[lo] + nums[hi]。
     • 如果 current_sum == 0，则找到了一个有效的三元组。将其添加到结果集中。然后，为了避免重复，需要移动 lo 和 hi 指针，跳过所有与当前 nums[lo] 和 nums[hi] 相同的元素。
     • 如果 current_sum < 0，说明和太小，需要增大。移动左指针 lo += 1。
     • 如果 current_sum > 0，说明和太大，需要减小。移动右指针 hi -= 1。
4. 去重处理: 在整个过程中，需要细致地处理重复元素，以确保最终结果集中不包含重复的三元组。
   • 外层循环去重: 当 i > 0 且 nums[i] == nums[i-1] 时，说明当前 nums[i] 与前一个元素相同，以它为第一个元素找到的三元组将是重复的，因此可以直接跳过。
   • 内层双指针去重: 当找到一个有效的三元组后，lo 和 hi 指针移动时，也需要跳过重复的元素。例如，while lo < hi and nums[lo] == nums[lo + 1]: lo += 1。

示例代码

以下是基于排序和双指针思想的优化版Python代码：

from typing import List

def threeSum(nums: List[int]) -> List[List[int]]:
    unique_triplets = []
    nums.sort()  # 1. 对数组进行排序

    # 遍历数组，固定第一个元素
    # 循环到 len(nums) - 2 是因为至少需要三个元素
    for i in range(len(nums) - 2):
        # 外层循环去重：如果当前元素与前一个元素相同，则跳过
        # 避免生成重复的三元组，例如 [-1, -1, 2] 和 [-1, -1, 2]
        if i > 0 and nums[i] == nums[i - 1]:
            continue

        # 使用双指针在剩余子数组中查找
        lo = i + 1  # 左指针从当前元素的下一个位置开始
        hi = len(nums) - 1 # 右指针从数组末尾开始

        while lo < hi:
            target_sum = nums[i] + nums[lo] + nums[hi]

            if target_sum < 0:
                lo += 1  # 和太小，左指针右移，增大和
            elif target_sum > 0:
                hi -= 1  # 和太大，右指针左移，减小和
            else: # target_sum == 0，找到一个三元组
                unique_triplets.append([nums[i], nums[lo], nums[hi]])

                # 内层双指针去重：跳过重复的lo元素
                # 避免生成重复的三元组，例如 [-2, 0, 2] 和 [-2, 0, 2]
                while lo < hi and nums[lo] == nums[lo + 1]:
                    lo += 1
                # 内层双指针去重：跳过重复的hi元素
                while lo < hi and nums[hi] == nums[hi - 1]:
                    hi -= 1

                # 找到一个有效三元组后，左右指针同时向内移动，继续寻找
                lo += 1
                hi -= 1
    return unique_triplets

时间复杂度分析

1. 排序: nums.sort() 的时间复杂度是 O(N log N)。
2. 主循环: 外层 for 循环迭代 N 次（len(nums) - 2）。
3. 双指针循环: 内层 while lo < hi 循环在最坏情况下，lo 和 hi 指针会遍历 N 次。例如，lo 从 i+1 移动到 hi，hi 从 len(nums)-1 移动到 lo，总共移动的步数是 O(N)。

因此，总的时间复杂度为 O(N log N + N * N) = O(N^2)。由于三数之和问题在没有额外数据结构辅助的情况下，至少需要检查 O(N^2) 对组合，因此 O(N^2) 是目前已知的最优时间复杂度。

空间复杂度分析

除了存储结果列表 unique_triplets 所需的空间（最坏情况下 O(N^3)，但实际通常远小于此），该算法主要依赖于输入数组的排序。如果 Python 的 sort() 方法是原地排序（通常是这样），那么额外的空间复杂度为 O(1)。如果排序算法需要额外空间（例如，某些版本的归并排序），则空间复杂度可能为 O(N)。在大多数竞争性编程场景中，通常认为它是 O(1) 的额外空间复杂度（不计输出空间）。

总结与注意事项

• 排序是关键: 排序是实现 O(N^2) 解决方案和高效去重的基础。
• 双指针的效率: 双指针技术在已排序数组中查找满足特定条件的元素对时非常高效，将内层循环的复杂度从 O(N) 降低到 O(1)（相对于遍历），从而将总复杂度从 O(N^3) 降低到 O(N^2)。
• 彻底去重: 必须在三个层面考虑去重：
  1. 固定第一个元素 nums[i] 时，避免重复。
  2. 找到有效三元组后，在移动 lo 指针时，跳过重复的 nums[lo]。
  3. 找到有效三元组后，在移动 hi 指针时，跳过重复的 nums[hi]。
• 边界条件: 确保 i、lo、hi 指针的范围正确，特别是 len(nums) - 2 这样的边界，以避免索引越界。

通过掌握这种排序加双指针的模式，你不仅能高效解决三数之和问题，还能将这种思想应用于其他类似的 N 数之和问题。

以上就是Python高效解决LeetCode三数之和问题：从超时到O(N^2)优化实践的详细内容，更多请关注php中文网其它相关文章！