Two-sum and frequency counts in Python

Why this matters

Many programming tasks ask a question about pairs or counts. "How many times does each word appear?" "Which two numbers add up to 100?" "How many pairs of equal values are there?" The slow way to answer is to compare every item with every other item. With 10,000 numbers, that is about 50 million comparisons. The fast way uses a dict to remember what you have already seen, so each item is handled once.

This lesson builds two patterns. The first is a frequency count, a dict from each value to the number of times it appears. Python's Counter class from collections builds one in a single call. The second is the two-sum pattern, which finds two values that add up to a target by looking back at earlier values. The dicts and sets lesson of this module introduced both in short form. This lesson takes them one step at a time, with a worked example at each step.

Words you need first

A frequency is the number of times a value appears in a list. In the list 3, 1, 3, the frequency of 3 is 2.

A frequency count is a dict where each key is a value from the list and each value is its frequency. A Counter is a dict made for this job.

A pair is two items at two different positions in a list. The values in a pair can be equal, but the positions must differ. The list [4, 4] has one pair of equal values, and its positions are 0 and 1.

A one-pass solution reads the list once, from the first item to the last, with one loop. It does not need a second loop inside the first.

A brute-force solution checks every possibility. For pairs, it compares each item with each later item. This takes n times (n minus 1) divided by 2 comparisons for n items.

Counting with a Counter

Start with the simplest case. The list is [3, 1, 3, 2, 3, 1]. You want to know how many times each value appears.

The Counter reads the list and counts each value. Its method most_common returns the values with the highest counts, largest first. A plain dict does the same job with get and a line of code, and the loop version is worth knowing because it is the pattern you will write in other problems.

The output is:

[(1, 2), (2, 1), (3, 3)]
most frequent: (3, 3)
{3: 3, 1: 2, 2: 1}

Trace the manual loop by hand. The first value is 3, and manual.get(3, 0) returns 0 because 3 is not yet a key, so the dict gets 3: 1. The second value is 1, so the dict gets 1: 1. The third value is 3 again, and get returns the stored 1, so 3 goes to 2. After all six values, the counts are 1 twice, 2 once and 3 three times.

The dict printed with 3 first, because a dict keeps the order in which keys were first added. The sorted(counts.items()) call prints the pairs in key order. Sort whenever you need a fixed order for output.

The get(v, 0) call is the important part. It returns 0 for a value that has not been seen yet, so the first count starts at 1. Reading the count with counts[v] would raise KeyError for a new value.

Two-sum: find two values that add up to a target

The two-sum problem asks: given a list of numbers and a target, find the positions of two numbers that add up to the target. Return the two positions, in any order.

The brute-force answer checks every pair. For the list [2, 7, 11, 15] and target 9, it tries 2 with 7, 2 with 11, and so on. That works, but it takes time that grows with the square of the list size.

Take a small case to see the difference. A list of 10 numbers has 45 pairs to check in brute force. The dict method does 10 lookups and 10 stores, and it stops early when it finds an answer. The gap grows fast: at 10,000 numbers, brute force checks about 50 million pairs, while the dict method does about 20,000 dict operations.

A dict makes it one pass. Walk through the list. For each number x, compute the partner it needs: target - x. Look in the dict to see whether the partner has already appeared. If it has, you have a pair. If not, store x and its position, and move on.

The dict stores value to position. The key is the number, and the value is the index where it was seen.

A worked example by hand

Take nums = [2, 7, 11, 15] and target = 9. Each row is one step of the loop.

The second row finds the pair. The partner 2 was stored at position 0, and the current number is 7 at position 1. So the answer is [0, 1]. Note that 11 and 15 are never looked at, because the loop stopped.

Here is the code with the same input, plus two more inputs:

The output is:

[2, 7, 11, 15], 9 -> [0, 1]
[3, 2, 4], 6 -> [1, 2]
[1, 2, 3], 100 -> None

The function enumerate gives the position and the value together, so the loop does not need to index the list. The second input shows that the partner can come from the middle of the list. For [3, 2, 4] and target 6: at position 0, the value 3 needs 3, which is not yet seen, so 3 is stored. At position 1, the value 2 needs 4, which is not seen, so 2 is stored. At position 2, the value 4 needs 2, which is in the dict at position 1. The answer is [1, 2].

The third input has no pair, so the loop finishes and the function returns None. A caller must check for None before using the result. The mistakes section below shows what happens without that check.

Why look up before storing

The order inside the loop matters. The code first checks whether the partner is in the dict, and only then stores the current number. This stops a number from pairing with itself. A single number x can match only an earlier position. The mistakes section below shows what goes wrong when the order is reversed.

Counting pairs: every match, not just one

Sometimes you do not need the positions, only how many pairs add up to the target. For example: how many pairs of positions i < j have nums[i] + nums[j] == target?

The same idea works. Keep a count of how many times each value has appeared so far. For each number x, the number of earlier values that complete a pair with x is the count of target - x. Add that count to the total. Then add x to the counts for the next steps.

A worked example: counting pairs

Take [1, 3, 3, 3, 5, 7] and target 6. Here is each step.

At position 3, the number 3 has two earlier 3s, so it makes two new pairs. The total is 4: the pair (1, 5), and three pairs of 3s. The pairs of 3s are the three ways to choose two of the three positions 1, 2 and 3.

This program does the same work, and it also checks a pair of equal values:

The output is:

[1, 3, 3, 3, 5, 7], 6 -> 4
[4, 4, 4, 4], 8 -> 6
[2, 8, 5], 10 -> 1

The second line is the interesting one. Four equal values give 6 pairs. The first 4 has no earlier partner, so it adds 0. The second 4 has one earlier 4, so it adds 1. The third adds 2, and the fourth adds 3. The total is 0 + 1 + 2 + 3 = 6. In general, n equal values give n times (n minus 1) divided by 2 pairs.

A second method for the same count

The same answer can come from the counts alone, without walking the list in order. Count every value first. Then look at each distinct value a and its partner b = target - a:

• If a is less than b, every copy of a pairs with every copy of b. The number of pairs is count(a) * count(b).
• If a equals b, the copies pair among themselves. With k copies, the number of pairs is k * (k - 1) // 2.
• Only handle each pair of values once. The rule a < b makes sure of that.

The output is:

[1, 3, 3, 3, 5, 7], 6 -> 4
[4, 4, 4, 4], 8 -> 6
[2, 8, 5], 10 -> 1

Both methods give the same answers. The one-pass method is the one to remember, because it needs no second pass over the counts. The frequency method is useful when the counts are already available, for example from an earlier step.

The operator // is integer division. It keeps the result a whole number, so k * (k - 1) // 2 is exact.

Python integers have no size limit, so the total never overflows. In languages with fixed-size integers, the total of pairs can wrap around, and the code must use a larger type. Python does not have that problem. Here is a list of 100,000 ones with a target of 2, which has about five billion pairs:

The output is:

pairs: 4999950000

The answer is exact. A 32-bit integer holds numbers only up to about 2.1 billion, so the same total would wrap around in a language that uses 32-bit integers for it.

Common mistakes

Storing before looking up. This is the most common two-sum bug. If the current number is stored first, the lookup can find the current number itself. Here is the wrong version, with the input [4] and target 8:

The output is:

[0, 0]

The answer uses position 0 twice. The number 4 was stored, then found as its own partner, because 8 minus 4 is 4. The fix is to check the dict before storing, as the correct version does.

Adding to a missing count. A dict raises KeyError when you read a missing key with []. Using get returns None instead, and adding 1 to None fails. The next program reads the count for a key that is not there:

counts = {"a": 1}
n = counts.get("b") + 1
print(n)

The output is an error:

TypeError: unsupported operand type(s) for +: 'NoneType' and 'int'

The message names the two types, NoneType and int, so you can see that the value came from a missing key. The fix is counts.get("b", 0) + 1, which uses 0 as the default. A Counter gives 0 for any missing key without a default, so Counter(...)["b"] returns 0.

Expecting None to be a list. The function two_sum returns None when no pair exists. A caller that does result[0] without a check gets TypeError: 'NoneType' object is not subscriptable. Check the result first with if result is None:, or return an empty list and say so in the docstring.

Sorting the list and losing the positions. Some other methods sort the list first, which is a different technique. The sort changes the positions, so the answer no longer says where the numbers were in the original input. If the task asks for positions, use the dict method in this lesson.

Practice

1. Count the letters in "banana" with a dict. Print the counts with the keys in sorted order. The output should be {'a': 3, 'b': 1, 'n': 2}.

   Show the answer:

   Copy

   counts = {}
   for ch in "banana":
       counts[ch] = counts.get(ch, 0) + 1
   print(dict(sorted(counts.items())))

   The printed line is {'a': 3, 'b': 1, 'n': 2}. The letter a appears three times, b once and n twice.

2. Given nums = [5, 75, 25] and target = 100, return the positions of the two numbers that add up to 100, using one pass and a dict.

   Show the answer: step through the loop. At position 0, the number 5 needs 95, which is not seen, so store 5 at position 0. At position 1, the number 75 needs 25, which is not seen, so store 75 at position 1. At position 2, the number 25 needs 75, which is in the dict at position 1. The answer is [1, 2].

3. Count the pairs of positions i < j in [2, 2, 3] where the two values add up to 4. Use the one-pass method.

   Show the answer: the only pair with sum 4 is the two 2s, at positions 0 and 1. The pair 2 and 3 sums to 5, so it does not count. The method adds 0 at position 0, adds 1 at position 1 (one earlier 2 is in the dict), and adds 0 at position 2 (it needs 1, which is not seen). The total is 1.

4. Count the pairs of positions i < j where the two values differ by exactly k, with k > 0. For example, with [1, 5, 3, 4, 2] and k = 2, the pairs are (1, 3), (3, 5) and (2, 4), so the answer is 3. Use one pass and a dict of counts.

   Show the answer: for each number x, the earlier numbers that differ from it by k are x + k and x - k. Because k is positive, these two values are different, so they never double-count. Add the counts of both, then add x to the dict:

   Copy

   def count_diff_pairs(nums, k):
       count = {}
       pairs = 0
       for x in nums:
           pairs += count.get(x + k, 0)
           pairs += count.get(x - k, 0)
           count[x] = count.get(x, 0) + 1
       return pairs

   The function returns 3 for the example. Each pair is found when the later number is read.

5. Write a function first_repeat that returns the first value that appears a second time when you scan the list from left to right. Return -1 if no value repeats. Use a set, and explain why checking membership before adding is enough to detect a repeat.

   Show the answer:

   Copy

   def first_repeat(nums):
       seen = set()
       for x in nums:
           if x in seen:
               return x
           seen.add(x)
       return -1

   The set answers "have I seen this?" in O(1) on average, which means about one step however long the list is. Checking before adding is enough, because a value is in the set only if an earlier position held it. The first value that is already in the set is the first repeat. For [2, 1, 3, 1, 2], the function returns 1, because the second 1 at position 3 comes before the second 2 at position 4. For [1, 2, 3], the function returns -1.

Summary

• A frequency count stores each value and its count. Counter(values) builds one, and get(v, 0) adds one to a count that starts at zero when the value is new.
• The one-pass two-sum method checks each number against the dict of earlier numbers before it stores the current one.
• For counting pairs, add the number of earlier matches to the total, then store the current value.
• Check the dict before you store the current value, or a number can pair with itself.
• Read a missing count with get(key, 0) or a Counter, not with [], and check None before you use a result.