python - 三和优化解

Question

我正在尝试解决如下所示的 3 Sum 问题:

Given an array S of n integers, are there elements a, b, c in S such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note: The solution set must not contain duplicate triplets.

这是我对这个问题的解决方案:

def threeSum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    nums.sort()
    n = len(nums)
    solutions = []
    for i, num in enumerate(nums):
        if i > n - 3:
            break

        left, right = i+1, n-1
        while left < right:
            s = num + nums[left] + nums[right] # check if current sum is 0
            if s == 0:
                new_solution = [num, nums[left], nums[right]]
                # add to the solution set only if this triplet is unique
                
                if new_solution not in solutions: 
                    solutions.append(new_solution)
                right -= 1
                left += 1
            elif s > 0:
                right -= 1
            else:
                left += 1

    return solutions

此解决方案在 O(n**2 + k) 的时间复杂度和 O(k) 的空间复杂度下运行良好，其中 n 是输入数组，k为解数。

在 LeetCode 上运行此代码时，我收到大型数组的超时错误。我想知道如何进一步优化我的代码以通过判断。

P.S:我已经阅读了this related question中的讨论。 .这并没有帮助我解决问题。

Answer 1

您可以对算法进行一些改进:

1) 使用sets而不是您的解决方案列表。使用集合将确保您没有任何重复项，并且您不必执行 if new_solution not in solutions: 检查。

2) 为全零列表添加边缘情况检查。开销不大，但在某些情况下可以节省大量时间。

3) 将枚举更改为第二个 while。它有点快。奇怪的是，我在使用 while 循环然后使用 n_max = n -2; 的测试中获得了更好的性能。 for i in range(0, n_max): 阅读 this xrange 或 range 的问答应该更快。

注意:如果我运行测试 5 次，我将不会获得相同的时间。我所有的测试都是+-100毫秒。因此，请对一些小的优化持保留态度。对于所有 python 程序，它们可能并不是真的更快。对于运行测试的确切硬件/软件配置，它们可能只会更快。

另外:如果你从代码中删除所有注释，它会快很多哈哈哈，比如快 300 毫秒。只是运行测试的一个有趣的副作用。

我已将 O() 符号放入您代码中需要花费大量时间的所有部分。

def threeSum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    # timsort: O(nlogn)
    nums.sort()

    # Stored val: Really fast
    n = len(nums)

    # Memory alloc: Fast
    solutions = []

    # O(n) for enumerate
    for i, num in enumerate(nums):
        if i > n - 3:
            break

        left, right = i+1, n-1

        # O(1/2k) where k is n-i? Not 100% sure about this one
        while left < right:
            s = num + nums[left] + nums[right]  # check if current sum is 0
            if s == 0:
                new_solution = [num, nums[left], nums[right]]
                # add to the solution set only if this triplet is unique

                # O(n) for not in
                if new_solution not in solutions:
                    solutions.append(new_solution)
                right -= 1
                left += 1
            elif s > 0:
                right -= 1
            else:
                left += 1

    return solutions

这是一些不会超时且速度很快的代码。它还暗示了一种使算法更快的方法(使用更多集合；))

class Solution(object):
def threeSum(self, nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    # timsort: O(nlogn)
    nums.sort()

    # Stored val: Really fast
    n = len(nums)

    # Hash table
    solutions = set()

    # O(n): hash tables are really fast :)
    unique_set = set(nums)
    # covers a lot of edge cases with 2 memory lookups and 1 hash so it's worth the time
    if len(unique_set) == 1 and 0 in unique_set and len(nums) > 2:
        return [[0, 0, 0]]

    # O(n) but a little faster than enumerate.
    i = 0
    while i < n - 2:
        num = nums[i]

        left = i + 1
        right = n - 1

        # O(1/2k) where k is n-i? Not 100% sure about this one
        while left < right:
            # I think its worth the memory alloc for the vars to not have to hit the list index twice. Not sure
            # how much faster it really is. Might save two lookups per cycle. 
            left_num = nums[left]
            right_num = nums[right]
            s = num + left_num + right_num  # check if current sum is 0
            if s == 0:
                # add to the solution set only if this triplet is unique
                # Hash lookup
                solutions.add(tuple([right_num, num, left_num]))
                right -= 1
                left += 1
            elif s > 0:
                right -= 1
            else:
                left += 1
        i += 1

    return list(solutions)