Two Sum – Python

I’ve been reviewing technical interview practice over the last few weeks, and thought it might be a good idea to consolidate my notes and think about how I would teach it. This always worked very well for music, and seems to work for writing code too.

I’m going to go over the Two Sum problem and use the two-pointer method. And I’m doing this in Python, but might add Ruby for it later. I like to see them next to each other to compare. Actually, I might do it in C# too and put them all next to each other. That’s what I do when I’m taking notes on my notion page, put the languages I’m working on next to each other. Anyway, off we go

Here is a description of the problem from leet code (linked above)

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.

I’m going to use pytest to write some simple tests first, and then solve it.

I want to test for the following

1. The solution exists
2. The solution doesn’t exist
3. What happens with negative numbers
4. What happens with duplicate numbers
5. Test an empty input
6. Test a single element being input

I think that’s pretty thorough, but let me know if I forgot anything

I’m doing all of this in one file, so this is at the bottom, you will have to run pip install pytest if you haven’t done so before

import pytest

def test_two_sum():
    # solution exists
    nums = [2, 7, 11, 15]
    target = 9
    assert two_sum(nums, target) == [0, 1]

    # solution doesn't exist
    nums = [2, 7, 11, 15]
    target = 10
    assert two_sum(nums, target) == []

    # negative numbers
    nums = [-1, 0, 1, 2, -1, -4]
    target = -1
    assert two_sum(nums, target) == [0, 2]

    # duplicate numbers
    nums = [3, 3]
    target = 6
    assert two_sum(nums, target) == [0, 1]

    # empty input
    nums = []
    target = 5
    assert two_sum(nums, target) == []

    # single element input
    nums = [5]
    target = 5
    assert two_sum(nums, target) == []

So there are a bunch of tests, and this is our end goal. Now lets try to solve them

I’m going to define the function two_sum and give it two parameters, nums which is the list of integers, and target which is the sum I want to find

I’m going to use an empty dictionary to make this more efficient. There is a brute force way to do it, but I am going for efficiency and this seems like it will be more efficient.

The dictionary is going to be num_indices and it will store the corresponding indices of every number that you find in the list, nums. You get constant time lookup in a dictionary, which is why I’m looking at it.

Here’s what I have so far

def two_sum(nums, target):
    num_indices = {}

I need to iterate over the list now, so I am going to use i to keep track of the index of each number in nums. I am also going to use enumerate to access the index and the value of every element.

I need a way to figure out if the target is attainable, so I will use total = target - num

The way this works is that I need to find the difference needed to actually get the target is 7 and num is 5, then the total is 2, and it is going to check if the dictionary has a key of 2.

If the total is already a key in the indices, it is going to return the index of num and the index of the key it found. If it doesn’t work, then it will take num and set it as a key with the value being it’s index, and move onto the next one. It looks like this

for i, num in enumerate(nums):
      total = target - num
      if total in num_indices:
          return [num_indices[total], i]
      num_indices[num] = i

And then when it all get’s put together, you have this

def two_sum(nums, target):
    num_indices = {}
    for i, num in enumerate(nums):
        total = target - num
        if total in num_indices:
            return [num_indices[total], i]       
        num_indices[num] = i
    return []

This should pass the tests I wrote earlier (make sure you have import pytest at the top)

You can run it with pytest two_sum.py

The entire block looks like this, with a test I added to put the numbers of the first test in a different order.

import pytest
def two_sum(nums, target):
    num_indices = {}
    for i, num in enumerate(nums):
        total = target - num
        if total in num_indices:
            return [num_indices[total], i]       
        num_indices[num] = i
    return []

def test_two_sum():
    nums = [2, 7, 11, 15]
    target = 9
    assert two_sum(nums, target) == [0, 1]
    
    # solution exists in different order
    nums = [7, 11, 15, 2]
    target = 9
    assert two_sum(nums, target) == [0, 3]


    nums = [2, 7, 11, 15]
    target = 10
    assert two_sum(nums, target) == []

    nums = [-1, 0, 1, 2, -1, -4]
    target = -1
    assert two_sum(nums, target) == [0, 1]

    nums = [3, 3]
    target = 6
    assert two_sum(nums, target) == [0, 1]

    nums = []
    target = 5
    assert two_sum(nums, target) == []

    nums = [5]
    target = 5
    assert two_sum(nums, target) == []

Alright, let me know if I missed anything, or if there is a better way!