Two Sum
Given an array of integers and a target value, return the positions of the two entries that add up to that target.
Looking for two entries that sum to 9. The map starts empty.
1def two_sum(nums, target):2 seen = {}3 for i, n in enumerate(nums):4 need = target - n5 if need in seen:6 return [seen[need], i]7 seen[n] = i8 return []Read the 11 steps as text
- 1Looking for two entries that sum to 9. The map starts empty.
- 2At index 0 the value is 3, so its partner would have to be 6.
- 3No 6 recorded yet, so store 3 at index 0 and carry on.
- 4At index 1 the value is 11, so its partner would have to be -2.
- 5No -2 recorded yet, so store 11 at index 1 and carry on.
- 6At index 2 the value is 15, so its partner would have to be -6.
- 7No -6 recorded yet, so store 15 at index 2 and carry on.
- 8At index 3 the value is 2, so its partner would have to be 7.
- 9No 7 recorded yet, so store 2 at index 3 and carry on.
- 10At index 4 the value is 7, so its partner would have to be 2.
- 112 is in the map at index 3. 2 + 7 = 9 โ that's the pair.
The idea
The brute force is to try every pair, which is O(nยฒ) and works fine โ it's just quadratic because it keeps re-asking a question it has already answered many times over.
Flip the question around. Instead of "which two numbers add to the target", ask, for each number as you pass it: "what would its partner have to be, and have I walked past that partner already?" The partner is fully determined โ it's target minus the current number โ so there is nothing to search for. There is only something to look up.
That turns the inner loop into a hash lookup, and a hash lookup is O(1). One pass is now enough, because by the time you reach the second element of the answering pair, the first one is already in the map.
The approach
- 1Keep a map from value to the index it was seen at.
- 2For each element, compute what it needs: need = target โ nums[i].
- 3If need is already a key in the map, you've found the pair โ return the stored index and the current one, in that order.
- 4Otherwise record nums[i] โ i and move on.
Complexity
One pass, and each step does O(1) expected hash work. The map holds at most n entries.
Code
def two_sum(nums, target):
seen = {}
for i, n in enumerate(nums):
need = target - n
if need in seen:
return [seen[need], i]
seen[n] = i
return []What goes wrong
- โขAdding the current number to the map before checking for its partner. If the target is exactly twice the current value, the element will match itself and you'll return the same index twice. Check first, insert second โ the order in the loop is the whole correctness argument.
- โขAssuming the array is sorted and reaching for two pointers. Two pointers is the better answer once it's sorted, but sorting costs O(n log n) and scrambles the original indices you were asked to return.
- โขReturning the values instead of the indices. It's the single most common misread of this problem.