Two sum is my favorite warm up question for a coding interview. Its easy to explain and easy to solve and there aren’t too many clarifying questions that need to be asked. Its one that you can just jump into.

The challenge

Given an array of integers and an integer target, return indices of the two numbers that sum to the target. Each input shall have but one solution.

Examples

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

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

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

The Solution

Our goal is to determine if there are two numbers that sum to a target. If we iterate over each value in the array we can derive the complimentary number by subtracting it from the target. If our target is 9 and our current value is 2, then the other number we’re looking for is 7. Now we just need a quick way to determine if the compliment is present in the array and, if so, its index.

We can utilize a hash map to store all values in the array and perform quick O(1) lookups. We can complete the solution in a single pass. If we miss the first value of the solution we’ll store it in the hash and retrieve it when we come upon its complement later in the array.

import pytest


class Solution:

    # Runs at O(n) time and O(n) space
    def twoSum(self, nums, target):
        hash_nums = {}
        for i, num in enumerate(nums):
            compliment = target - num
            if compliment in hash_nums:
                return [hash_nums[compliment], i]
            hash_nums[num] = i

@pytest.mark.parametrize("nums,target,expected", [
    ([3,2,4], 6, [1,2]),
    ([2, 7, 11, 15], 9, [0, 1])
])
class TestClass:
    def test_simple_case(self, nums, target, expected):
        s = Solution()
        assert s.twoSum(nums, target) == expected