[python 刷题] 167 Two Sum II - Input Array Is Sorted & 15 3Sum

[python 刷题] 167 Two Sum II - Input Array Is Sorted & 15 3Sum

虽然 3 sum 出来的比较早，不过按照解法来说，2 sum II 算是 3 sum 的前置解法

167 Two Sum II - Input Array Is Sorted

题目：

Given a 1-indexed array of integers numbers that is already sorted in non-decreasing order, find two numbers such that they add up to a specific target number. Let these two numbers be numbers[index1] and numbers[index2] where 1 <= index1 < index2 < numbers.length.

Return the indices of the two numbers, index1 and index2, added by one as an integer array [index1, index2] of length 2.

The tests are generated such that there is exactly one solution. You may not use the same element twice.

Your solution must use only constant extra space.

2Sum II 的前置肯定就是 2 Sum 了，不过这题比较难的地方在于空间复杂度只能为 $O(1)$，也就是说，不能用额外的 dict 去存储补数

这题的解法就是双指针：

想像一下左右各有一个指针，如果想要两数之和变大，那么左边的指针可以向右移动。如果想要两数之和变小，那么右边的指针可以向左移动

在数组中没有解的情况下，一旦双指针相遇，这个循环就可以结束了，因此最差情况下就是遍历整个数组，所以时间复杂度为 $O(n)$。只需要保留两个指针，空间复杂度为 $O(1)$

解法如下：

class Solution:
    def twoSum(self, numbers: List[int], target: int) -> List[int]:
        l, r = 0, len(numbers) - 1
        while l < r:
            if numbers[l] + numbers[r] == target:
                return [l + 1, r + 1]

            if numbers[l] + numbers[r] < target:
                l += 1
            else:
                r -= 1

15 3Sum

题目：

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.

Notice that the solution set must not contain duplicate triplets.

其实这题可以理解成一个嵌套版本的 2sum，因此先决条件就是进行一下排序。以题目中给的案例 [-1,0,1,2,-1,-4] 来说，排序完的结果为：[-4, -1, -1, 0, 1, 2]，这个情况下，就可以把 3 sum 拆解成 2 sum II。

即 target 为 nums[i]，找出 nums[j] + nums[k] = -nums[i] 的搭配，图解的话，大概就是这么理解的:

代码如下：

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        res = collections.defaultdict(set)
        nums.sort()

        i = 0
        while i < len(nums) - 2:
            val = nums[i]
            j, k = i + 1, len(nums) - 1
            while j < k:
                if nums[j] + nums[k] + val == 0:
                    res[val].add((nums[j], nums[k]))
                    j += 1

                if nums[j] + nums[k] + val > 0:
                    k -= 1
                else:
                    j += 1

            i += 1

        return [[key, x, y] for key, val in res.items() for x, y in val]

空间复杂度为 $O(n)$，首先排序有可能会创建一个新的数组，其次这里的确用了额外的 dict 去保存 tuple 的 set（为了不重复）

时间复杂度是 $O(n^2)$，虽然排序是 $O(nlog(n))$ 的复杂度，不过下面两个循环爆了排序，成了 dominant factor

返回值用了 python 的 list comprehension 技巧，是一个可以把两个循环压缩成一行，并且生成 list 的方法。python 内部对 list comprehension 的实现有优化，一般来说 list comprehension 的效率比循环快

这行代码解析起来就是：

• 从 res.items() 中获取 key（数字）, val（tuple） 的 k-v 对
• 从 tuple 中获取 x 和 y（对应的下标）
• 遍历将将 key, x, y 组成一个数组，形成一个数组的数组

这是没做什么优化的解，不过另一个优化的方式如下：

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        res = []
        nums.sort()

        for i, a in enumerate(nums):
            # Skip positive integers
            if a > 0:
                break

            if i > 0 and a == nums[i - 1]:
                continue

            l, r = i + 1, len(nums) - 1
            while l < r:
                threeSum = a + nums[l] + nums[r]
                if threeSum > 0:
                    r -= 1
                elif threeSum < 0:
                    l += 1
                else:
                    res.append([a, nums[l], nums[r]])
                    l += 1
                    r -= 1
                    while nums[l] == nums[l - 1] and l < r:
                        l += 1

        return res

• enumerate() 会提供给数组的下标和当前值

  当前值为 nums[i]

• 当 nums[i]>0 时，终止循环

  这是一个小优化，题目的需求本来就是 nums[i] + nums[j] + nums[k] == 0，这里的数组又进行了排序，当 nums[i]>0 时，nums[j] 和 nums[k] 必然也 >0

• if i > 0 and a == nums[i - 1]: 这里去除重复值

  对比未被优化的代码，也就是跳过 nums[i] 已经作为 key 在 dict 的情况

• l & r 的双指针实现一致

• 最后的 while 也是另一个优化

  也是去除 tuple 中可能会出现重复值的方法，如：[1,1,1,2,2]，这里会出现的所有 combo 为 (1, 1), (1, 2), (2, 2)

  如果 target 是 2，那么右边的指针就会一直减少，一直到指向的区域只有 1，并且获取第一个解

  这个时候只要维持右边的指针不动，左边在指针重复的情况下一直向右移动，那么当遇到 l==r 的情况下，循环也会终止

  所以这里只需要移动一边就行了

总体的时间和空间复杂度依旧维持不变，不过一些小的优化也能够提速不少