1. Two Sum

Problem

Given an array of integers nums and an integer target, return indices of the two numbers such that they add up to target.

You may assume that each input would have exactly one solution, and you may not use the same element twice.

You can return the answer in any order.

https://leetcode.com/problems/two-sum/

Example 1:

Input: nums = [2,7,11,15], target = 9
Output: [0,1]
Explanation: Because nums[0] + nums[1] == 9, we return [0, 1].

Example 2:

Input: nums = [3,2,4], target = 6
Output: [1,2]

Example 3:

Input: nums = [3,3], target = 6
Output: [0,1]

Constraints:

2 <= nums.length <= 10⁴
-10⁹ <= nums[i] <= 10⁹
-10⁹ <= target <= 10⁹
Only one valid answer exists.

Follow-up: Can you come up with an algorithm that is less than O(n²) time complexity?

Test Cases

1 2	class Solution: def twoSum(self, nums: List[int], target: int) -> List[int]:

solution_test.py

import pytest

from solution import Solution
from solution2 import Solution as Solution2


@pytest.mark.parametrize('nums, target, expected', [
    ([2,7,11,15], 9, [0, 1]),
    ([3,2,4], 6, [1, 2]),
    ([3,3], 6, [0, 1]),
])
class Test:
    def test_solution(self, nums, target, expected):
        sol = Solution()
        result = sol.twoSum(nums, target)
        assert sorted(result) == sorted(expected)

    def test_solution2(self, nums, target, expected):
        sol = Solution2()
        result = sol.twoSum(nums, target)
        assert sorted(result) == sorted(expected)

Thoughts

一方面是要找到那两个相加等于目标值的数，另一方面需要能记录到这两数的原始索引下标。

用 O(n log n) 时间对数组排序。

从最小的数字开始遍历，直到 target / 2（包含）。

对于一个数字 v，用二分法，在后半数组中查找 target - v。二分查找是 O(log n)，总共也是 O(n log n)。

需要记录数字在排序前的数组下标。

Code

solution.py

from typing import List


class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        n = len(nums)
        indices = list(range(n))
        indices.sort(key=lambda i: nums[i])

        for i, p1 in enumerate(indices):
            v2 = target - nums[p1]
            p2 = self._bin_find(nums, indices, v2, i + 1)
            if p2 is not None:
                return [p1, p2]

        return []

    def _bin_find(self, nums: List[int], indices: List[int], value: int, left: int) -> int | None:
        l = left
        if l >= len(indices) or (t := nums[indices[l]]) > value:
            return None
        elif t == value:
            return indices[l]

        r = len(indices) - 1
        if (t := nums[indices[r]]) < value:
            return None
        elif t == value:
            return indices[r]

        while l <= r:
            m = (l + r) >> 1
            if (t := nums[indices[m]]) == value:
                return indices[m]
            elif t < value:
                l = m + 1
            else:
                r = m - 1

        return None

快一些

可以利用哈希表 O(1) 查询时间的特点，把所有数字放进哈希表，然后直接利用哈希表查找 target - v 是否存在。

因为题目限定了有唯一解，所以不用管重复的数字。唯一需要考虑的是，当 target 是偶数时，可能有两个 target / 2。

时间复杂度为 O(n)。

solution2.py

from typing import List


class Solution:
    def twoSum(self, nums: List[int], target: int) -> List[int]:
        indices = {}
        half = target >> 1 if target & 1 == 0 else None
        for i, v in enumerate(nums):
            # Special treatment of `target / 2` when target is even.
            if half is not None and v == half and v in indices:
                return [indices[v], i]

            indices[v] = i

        for i, v1 in enumerate(nums):
            if half is not None and v1 == half:
                continue

            v2 = target - v1
            if v2 in indices:
                return i, indices[v2]

        return []

本文除了题目描述（Problem）部分之外，其他内容由 Calf 原创，采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处。