Embercode
✕
Dicts and SetsLesson 3 of 17

55 min read

Two-sum, frequency counting and subarray sums in Python

Why it matters

Most interview problems that use a dict or a set share one move. Instead of looking back through the data each time you need an answer, you store what you have seen so far in a dict or set, and you answer each question in one step. That move turns many O(n^2) solutions into O(n) solutions. It is one of the most useful habits you can build, and it appears in problems about lists, strings, streams of events and graphs.

This lesson teaches the pattern through six concrete problems. For each one you will see the slow version first, then the dict or set version, a trace of the steps, the cost, and the mistakes people make. Every program is complete, so you can run it, change the numbers and predict the output.

The one move: remember what you have seen

Picture a door with a guest list. A person arrives, the guard checks the list to see if they are already inside, and then writes their name down. The guard does not walk back through the room to check every guest. The list answers "have I seen this person?" in one glance. A set does the same job for values, and a dict does it for values with extra information attached, such as an index or a count.

The pattern has two parts that must happen in the right order. For each item you ask the dict a question, and then you record the item. The order matters, and you will see why in the two-sum section. Many bugs in this lesson come from swapping the two steps.

Before you write any code, ask three questions about the problem. What is the question for each item? What must you remember from the earlier items? Does the order of the items matter? If you can answer those, the dict usually follows directly.

Pattern 1: frequency counting

Frequency counting means counting how often each item appears. It solves a long list of problems: finding the most common word, checking whether two strings are anagrams, finding the first character that appears only once, and deciding whether a number appears more than half the time.

The standard tool is collections.Counter. Give it any iterable, and it counts each item. A Counter is a dict, so the usual dict methods work on it, and a missing key counts as 0. The program below counts letters, finds the first letter that appears once, and shows a faster method for lowercase letters only.

Python · editable
{'a': 3, 'b': 1, 'n': 2}
first unique in banana: b
first unique in aabbcc: -
i count: 4

Start with first_unique. It must find the first letter in the original text that appears once, so the order in which letters were first seen matters. A dict keeps insertion order, and Counter is a dict, so count.items() walks the letters in the order they first appeared. For banana the counts are b=1, a=3 and n=2, so b is the first letter with count 1. For aabbcc no letter has count 1, so the function returns the placeholder -.

The dict(sorted(...)) line is only for printing. It puts the letters in alphabetical order so the output is easy to read.

The last part uses a list of 26 integers. The function ord returns a character's code, so ord("a") is 97 and ord(ch) - ord("a") gives an index from 0 to 25. When the characters are known to be lowercase letters, a list is faster than a dict and uses less memory. Use the dict when the characters can be anything, such as capital letters, digits or punctuation.

Trace the list code for mississippi. The letter i appears at four positions, so letters[ord("i") - ord("a")] is letters[8], and it holds 4.

The cost is O(n) time for a string of length n. The space is O(k) for k distinct characters with a dict, or O(26) with the list.

Pattern 2: two-sum

The two-sum problem is the classic first step into this pattern. You get a list of numbers and a target, and you return the indexes of two numbers that add up to the target. Assume exactly one answer exists.

The slow method checks every pair. It is simple, and you should write it first in an interview to show you understand the problem.

def two_sum_slow(nums, target):
    for i in range(len(nums)):
        for j in range(i + 1, len(nums)):
            if nums[i] + nums[j] == target:
                return [i, j]
    return [-1, -1]

The inner loop runs about n times for each of the n outer steps, so the time is O(n^2). The extra space is O(1). For 100,000 numbers that is about 5 billion checks, which is too slow.

The fast method uses the pattern. Walk the list once, from left to right. For each number x, the partner you need is target - x. Ask the dict whether you have already seen that partner. If you have, you have found the pair. If not, record x with its index and continue. Because you check before you record, the number cannot pair with itself.

Trace [2, 7, 11, 15] with target 9. The dict starts empty.

IndexNumberPartner neededPartner seen?Action
027norecord 2 at index 0
172yes, at index 0return [0, 1]

The walk ends at index 1, so the rest of the list is never checked. The answer is [0, 1], because 2 + 7 is 9.

Here is the complete program. It also counts pairs rather than returning a single pair, and it handles the no-answer case.

Python · editable
0 1
1 2
-1 -1
pairs: 4

The function enumerate gives each number together with its index, so the loop does not need a counter. print(*two_sum(...)) unpacks the returned list into two values, which print shows with a space between them.

Trace the second call, two_sum([3, 2, 4], 6). At index 0 the number is 3 and the partner needed is 3. The dict is empty, so 3 is recorded at index 0. At index 1 the number is 2 and the partner needed is 4. The dict does not have 4, so 2 is recorded at index 1. At index 2 the number is 4 and the partner needed is 2. The dict has 2 at index 1, so the function returns [1, 2]. The numbers are 2 + 4 = 6. The first 3 was never paired with itself, because it was not in the dict when it was checked.

The third call has no answer, because 1 + 2 is 3, not 10. The loop finishes and returns the placeholder [-1, -1]. Always decide what the function returns when no answer exists, and say so in your code.

The count_pairs function asks a different question. It counts all pairs, not just one. The idea is the same. For each x, add the number of earlier items that equal target - x, then record x. The trace for [1, 5, 3, 3, 3] with target 6 is:

  • 1: no partner 5 yet, so 0 pairs. Record 1.
  • 5: partner 1 seen once, so 1 pair. Record 5.
  • 3: partner 3 not seen, 0 pairs. Record 3.
  • The second 3: partner 3 seen once, so 1 pair. Record another 3.
  • The third 3: partner 3 seen twice, so 2 pairs. Total 4.

The three 3s form three pairs, at positions 2 and 3, 2 and 4, and 3 and 4. The pair (1, 5) makes the fourth. The answer is 4.

In Java this count would need a long to avoid overflow. Python integers do not overflow, so the same code works for any size of answer. The section on objects below covers what that costs.

Two-sum on a sorted list: the two-pointer method

Your interviewer may ask, "What if the list is already sorted?" Then you have another tool. Put one pointer at the left end and one at the right end. Add the two numbers. If the sum is too small, move the left pointer right to get a bigger number. If the sum is too big, move the right pointer left. Stop when the sum matches or the pointers meet.

10left31426384115156right

sum: not computed yet

The left and right pointers are labelled under the boxes. Outlined boxes are the pair being added, filled boxes are the pair that works.

Find two items that add up to 14. One pointer starts at each end.

Step 1 of 4

Step through the visualizer above and watch the two pointers close in on each other. Every step moves one pointer, so the walk is O(n). The pointers use O(1) extra space, because the method needs no dict. The trade is that you must sort first, which costs O(n log n), and sorting moves the numbers, so you must remember their original indexes if the problem asks for positions. The dict method needs O(n) extra space but works on unsorted data in one pass. Choosing between them is a real design decision, and a good answer names both.

Pattern 3: anagram groups and checks

Two words are anagrams when they use the same letters the same number of times, such as listen and silent. Two common checks use the pattern.

To check one pair, count the letters of each word in one list of 26 integers, adding for the first word and subtracting for the second. If every count ends at zero, the words are anagrams. The check runs in O(n) time and uses constant space.

To group many words, you need a key that all anagrams share. Sorting the letters of a word gives a key that is the same for every anagram. listen and silent both sort to eilnst. Use that sorted form as the key in a dict, and add each word to the list for its key. The sorted function returns a list of characters, so "".join(...) turns it back into a string. The program below does both.

Python · editable
['eat', 'tea', 'ate']
['tan', 'nat']
['bat']
True False

The defaultdict(list) creates an empty list the first time a key is read, so the code does not need an if before each append. In group_anagrams, the first word eat sorts to aet, so it becomes the key for a new list. Then tea sorts to aet as well, and it joins the same list. tan sorts to ant, which is a new key. ate joins the aet list, and nat joins the ant list. bat sorts to abt, a third key. A plain dict keeps insertion order, so the groups come out in the order their first word appeared.

The cost is O(n * k log k) for n words of length at most k, because each word is sorted once. For short words that is effectively linear. If the words are long and many, you can use a count-based key instead, such as the 26 letter counts joined into a string, which costs O(n * k).

In is_anagram the counts go up for a and down for b. For listen and silent every letter goes up and comes back down, so every count is zero at the end. For rat and car the t count is left at 1 and the c count at -1, so the function returns False. A shorter version of the pair check is Counter(a) == Counter(b), which is easy to read but builds two dicts. The list version is the one interviewers expect.

Pattern 4: duplicates and the first repeat

Detecting a duplicate is the simplest use of a set. The question "have I seen this?" needs only in and add. The code is short, and the cost is O(n).

The next program finds the index of the first repeated value, which is a common follow-up. It returns the index at which the repeat happens, so the caller can also see where the first copy was. It then prints every duplicate name it sees in a list of names.

Python · editable
3
-1
duplicate: ana

In first_repeat, the list [2, 5, 1, 5, 2] gives these steps. Index 0 adds 2, index 1 adds 5, index 2 adds 1. Index 3 holds 5, and 5 is already in seen, so the function returns 3. The 2 at index 4 is also a repeat, but the function stops at the first one, which is the point of the question. The second call has no repeat, so the function returns -1.

The name loop prints duplicate: ana at the third name, because the set already holds ana. A check like this is the first step of many real tasks, such as finding a duplicate user id before saving a record.

Python also offers a shorter test for "are there any repeats": compare len(set(items)) < len(items). It is correct, and it is O(n). It does not tell you which value repeats first, so use the loop when you need the position.

Pattern 5: the longest run of consecutive values

Here is a problem that looks like it needs sorting. Given a list of numbers in any order, find the longest run of consecutive integers, such as 1, 2, 3, 4. Sorting solves it in O(n log n). The set method solves it in O(n).

Put every value in a set. Then walk the set, and for each value ask a single question: is this value the start of a run? It is a start if the value one lower is not in the set. Only starts are worth counting. From a start, count upward while the next value is present. Each value is counted from its start once, and no value is counted from the middle of a run. That is why the total work is linear.

Trace [100, 4, 200, 1, 3, 2]. The set holds those six values. The start 100 has no 99, so count upward: 101 is missing, so the run is 1 long. The value 4 has a 3, so it is not a start. The value 200 is a start with a run of 1. The value 1 has no 0, so it is a start. Count upward: 2, 3 and 4 are present and 5 is not, so the run is 1, 2, 3, 4, length 4. The values 3 and 2 have lower neighbours, so they are skipped. The answer is 4.

Python · editable
4
3
0

The second call holds 0, -1, -2 and 5. The start is -2, because -3 is missing, and the run -2, -1, 0 has length 3. The value 5 is a run of 1. The third call is an empty list, so the set is empty, the loop never runs, and the answer is 0.

The order in which a set gives its items does not matter here. The answer is the same for any order, which is why the loop can walk the set directly.

Why is this O(n)? Each value is checked against its lower neighbour once in the outer loop. The while loop runs only from a start, and the values it visits belong to that one run. A value in the middle of a run is never a start, so it never starts a count. Add up the work and each value is visited about twice. This is the kind of argument interviewers want to hear. Say it in your own words.

Pattern 6: subarray sums with prefix sums

Now the hardest pattern in this lesson. A subarray is a continuous slice of a list. The problem is to count the subarrays whose sum is exactly k. The numbers may be negative, and that rules out the sliding window technique you may know from other problems. A window grows when you add a positive number and shrinks when you remove one, but a negative number makes the sum go down when the window grows, so you cannot tell when to shrink. A different tool is needed.

The tool is the prefix sum. The prefix sum at position i is the total of all numbers from the start up to and including position i. A subarray from position j + 1 to position i has the sum prefix[i] - prefix[j], because the part from the start up to j is subtracted out. For that sum to equal k, you need prefix[i] - prefix[j] = k, which means prefix[j] = prefix[i] - k.

So for each i you ask: how many earlier prefix sums equal prefix[i] - k? Keep a dict from each prefix value to how many times it has appeared. Start the dict with 0 mapped to 1, which stands for the empty prefix before the first number. This starting entry handles subarrays that begin at the first position. Without it, a subarray such as [1, 2] with k = 3 would be missed, because its prefix sum of 3 would not find a partner.

Here is the trace for [1, 2, 3, -3, 3] with k = 3. The dict starts as {0: 1}.

PositionNumberPrefixNeed (prefix minus k)Times need seenRunning countDict after step
011-200{0:1, 1:1}
123011{0:1, 1:1, 3:1}
236312{0:1, 1:1, 3:1, 6:1}
3-33013{0:1, 1:1, 3:2, 6:1}
436325{0:1, 1:1, 3:2, 6:2}

Read position 1. The prefix is 3, and the need is 0. The dict holds 0 once, from the starting entry, so the subarray from position 0 to position 1 counts. That subarray is [1, 2], whose sum is 3. At position 2 the prefix is 6 and the need is 3. The dict holds 3 once, from position 1, so the subarray from position 2 to position 2 counts. That is [3]. At position 3 the number is -3, which brings the prefix back to 3. The need is 0, which the starting entry supplies, so the subarray [1, 2, 3, -3] counts. At position 4 the prefix is 6 again, and the need is 3. The dict holds 3 twice, from positions 1 and 3, so two subarrays end here: [3] at position 4, and [3, -3, 3] from position 2 to position 4. The total is 5.

Python · editable
fast = 5 slow = 5
random checks passed: 200

The program contains two functions that answer the same question. count_slow tries every start and every end, which is O(n^2). It is easy to trust, because it does exactly what the question says. count_fast is the prefix sum method, which is O(n) time and O(n) space. The test code checks that both agree on the example and on 200 random lists. This is a strong habit. When you write a clever function, write the slow one too, and compare them on small random inputs.

The generator is random.Random(7), which uses a fixed seed, so the same numbers appear every time you run the program on the same Python version. Python's random numbers are not the same as Java's for the same seed, so the lists differ from the Java version, but the check is the same. randrange(-5, 6) gives numbers from -5 to 5, and randrange(-3, 4) gives targets from -3 to 3. Negative numbers and negative targets are exactly the cases that break a sliding window, so they make a good test.

The answer for a list of 100,000 zeros with k = 0 is larger than 4 billion. count_fast([0] * 100000, 0) returns 5000050000, which Python stores exactly with no special type. In Java this answer needs long.

What happens inside: objects and memory

In Python every value is an object. A number in a dict is a full object with a type, a size and a reference count. Python does not need a separate boxed type for int, so you never write the wrapper code that Java needs. The cost is memory. On Python 3.14, a small int object takes 28 bytes, and 10**100 takes 72 bytes. The exact size depends on the Python version and the machine. A list of a million numbers holds a million objects, plus the list itself.

For most problems this cost is fine. For a problem with a very large input, you can reduce the cost. A list of 26 integers, as in the letter count, is much smaller than a dict of 26 keys. The array module from the standard library stores plain numbers without an object for each one. Use these tools only after you measure a real problem, because the readable version is usually fast enough.

The dict methods also matter for speed. The lines if need in seen: value = seen[need] look up the key twice. The line seen.get(need, 0) looks it up once. In a loop that runs a million times, that difference is noticeable. Prefer the single lookup form.

Choosing the pattern

Most problems in this lesson name their pattern in the question, once you know the signs. The table lists the signs. Read the question, find the row that matches, and then check the edge cases in the last column before you write code.

What the question asksPatternStructureWatch for
Two items that add up to a targetTwo-sumdict from value to indexThe item must not pair with itself
How often each item appearsFrequency countingCounter, or a list of 26 countsCase, punctuation and key type
Are two words anagramsLetter countslist of 26 with add and subtractLengths must match first
Group items that share a propertyGroupingdefaultdict(list) with a sorted keyKey must be the same for all group members
Is there a repeatSeen setset with in before addDecide what to return when there is none
Longest run of consecutive valuesRun from startset with start checkDuplicates and negative values
Count subarrays with a given sumPrefix sumsdict from prefix to countStarting entry 0 mapped to 1

The last column matters as much as the first. Most wrong answers in interviews are not wrong algorithms. They are correct algorithms that miss an edge case, such as a repeated value in two-sum or a zero target in prefix sums.

Explain the answer out loud

Interviewers judge how you think as much as the code you write. When you finish a solution, say three things in plain words. First, the idea: "I walk the list once and remember the numbers I have seen, so each partner is a lookup." Second, the cost: "That is O(n) time and O(n) space, instead of the O(n^2) pair check." Third, the test you will run: "I will try duplicates and a case with no answer, and check the result against the slow version."

This habit also helps you when you are stuck. If you cannot say the idea in one sentence, the idea is not clear yet. Go back to the table, find the pattern, and try again. A solution you can explain is far more useful to you, and to the person who reads your code later, than one that happens to pass.

Common mistakes

Recording before checking in two-sum. If you store seen[x] = i before you look for the partner, the number can pair with itself. With [3] and target 6, the code finds the partner at index 0 and returns [0, 0], which is wrong. Check first, then record. The output below shows the wrong version.

Python · editable
[0, 0]

Forgetting the starting entry in prefix sums. Without seen = {0: 1}, any subarray that starts at position 0 and has sum k is missed. The trace above shows the first subarray, [1, 2], counted only because the starting entry exists.

Using a sliding window with negative numbers. The window technique assumes that adding a number increases the sum. Negative numbers break that assumption. Use prefix sums when the list can hold negatives.

Putting a list inside a set. Sets and dict keys must be hashable, and a list is not. The error names the problem:

seen = set()
seen.add([3, 4])
TypeError: cannot use 'list' as a set element (unhashable type: 'list')

Use a tuple, such as seen.add((3, 4)), when the items are pairs.

Looking up a key that may be missing. seen[need] raises KeyError when the key is missing. Use seen.get(need) or need in seen first.

seen = {}
print(seen[4])
KeyError: 4

Forgetting what a no-answer returns. Two-sum needs a value for the case where no pair exists. Decide it up front and document it, such as returning [-1, -1], and test that case.

Debugging walkthrough: two-sum returns the wrong indexes

Suppose your two-sum function returns [1, 1] for an input that should give two different positions. This happens when the code pairs a number with itself, or when it returns the index of the first copy of a value by mistake. Work through these steps.

  1. Print the dict each iteration. Add print(i, nums[i], seen) inside the loop. You will see exactly which numbers are recorded and when.
  2. Look at the order. If the dict stores the current number before the check, you have the bug from the common mistakes section. Move the store to after the if need in seen check.
  3. Check the return values. The function returns [seen[need], i]. The first index is the earlier one. If the function returns them the other way around, the test may still pass, but the problem may ask for a specific order. Read the statement.
  4. Try a duplicate. Use [3, 3] with target 6. The correct answer is [0, 1]. If you get [0, 0], the code paired the first 3 with itself.
  5. Try a negative. Use [-2, 5] with target 3. The partner of 5 is -2, so the answer is [0, 1]. This checks that negative values work with target - x.

Each step tests one part of the logic, so the bug shows up quickly.

Edge cases

Empty lists give 0 for counts and -1 for "no answer". A single element can never form a pair, so two-sum returns the no-answer value. Duplicate values are the most important case in two-sum: the same value can appear twice, and the pair must use two different positions. Negative targets and negative numbers both work, because target - x is ordinary arithmetic.

For prefix sums, a target of 0 is special. Any run of zeros gives many subarrays with sum 0. The starting entry handles the subarrays that begin at the first position, and the dict handles the rest. Test a run of zeros to be sure, as the list of three zeros does in the practice section.

Large answers need no special care in Python, because integers grow as needed. The answer for 100,000 zeros with k = 0 is 5000050000, and Python prints it exactly.

Interview angle: state the pattern and the cost

When you meet a problem about pairs, sums, counts or repeats, say the pattern out loud. A good answer goes in three steps. First, state the brute force and its cost, such as O(n^2). Second, state the dict or set idea and its cost, such as O(n) time and O(n) space. Third, name the edge cases you will test, such as duplicates, negatives, empty input and no answer.

Ask clarifying questions before you code. Is the list sorted? Can values repeat? Do you need indexes or values? Is the answer unique? Can the sum be large? Each answer changes the solution, and asking shows the interviewer that you think about the problem before you write code.

Keep the cost table in mind. Two-sum with a dict is O(n) time and O(n) space. Two-sum on a sorted list with two pointers is O(n log n) time for the sort and O(1) extra space for the pointers. Frequency counting is O(n). Anagram grouping with sorted keys is O(n * k log k). The longest run is O(n). Subarray sums with prefix sums are O(n).

Real-world uses

Frequency counts drive search ranking, spam filters and word clouds. Two-sum-style lookups appear in matching systems, such as pairing a buy order with a sell order at a matching price. Duplicate detection protects against double charges and repeated uploads. Grouping anagrams is a close cousin of grouping records by a normalised key, such as lowercasing and trimming a customer's email before comparing. Prefix sums power range queries in analytics, such as the total sales between two dates, and they show up in image processing for computing sums over rectangles.

Practice

Try each exercise before you open its answer. Use paper for the traces, and check your answers with a short program.

  1. Check whether two strings are anagrams using a list of 26 integers. Hint: add for the first string, subtract for the second, and check that every count is zero.

Show the answer

Check the lengths first. Then loop over the two strings together with zip, add one to ord(x) - ord("a") and subtract one for ord(y) - ord("a"). If all(v == 0 for v in count) is true at the end, the strings are anagrams.

  1. Find the first repeated value in a list and return the value, not its index. Hint: a set and in.

Show the answer

Python · editable
4 None
  1. Trace count_fast for [1, 1, 1] with k = 2 and write the running counts. Hint: the answer is 2.

Show the answer

The dict starts as {0: 1}. The prefixes are 1, 2 and 3. At prefix 1 the need is -1, count stays 0. At prefix 2 the need is 0, found once, count 1. At prefix 3 the need is 1, found once, count 2. The answer is 2.

  1. Write a function that returns the number of distinct values in a list of integers. Hint: len of a set.

Show the answer

Python · editable
3
  1. Change two_sum so it returns every pair of indexes that adds up to the target, not just the first one. Hint: store a list of indexes for each value in the dict.

Show the answer

Python · editable
[[0, 1], [2, 3], [2, 4], [3, 4]]

The list is the same as the pair count from this lesson, shown as index pairs.

  1. Given a string, return the first character that appears more than once. Hint: a set, and return as soon as in is true.

Show the answer

Python · editable
r
  1. Write a function that checks whether a list contains two equal values that are at most k positions apart. Hint: a dict from value to its last index, and compare the distance.

Show the answer

Python · editable
True False
  1. Explain in your own words why the longest-run method visits each value only about twice. Hint: a start is a value whose lower neighbour is missing.

Show the answer

Each value is either a start, which is counted once from its start, or a middle value, which is skipped. Every value in a run is checked against its lower neighbour once, and only the start walks upward, so the total work is about twice the number of values.

  1. A teammate proposes a sliding window for "subarray sum equals k" on a list with negative numbers. Explain why it fails with a three-element example. Hint: try [3, -1, -1] with k = 1.

Show the answer

A window that shrinks when its sum is above k starts with [3], whose sum 3 is above 1. It drops the 3 at once, so the window never reaches [3, -1, -1], whose sum is 1. The rule to shrink only works when adding numbers never lowers the sum, and negative numbers break that.

  1. Count the subarrays whose sum is exactly 0 for [1, -1, 2, -2]. Use the prefix sum method and write the dict after each step. Hint: the answer is 3, and the prefixes are 1, 0, 2, 0 with a starting 0.

Show the answer

Start with {0: 1} and a count of 0.

  • Number 1: prefix 1, need 1, not in the dict, count stays 0. Dict becomes {0:1, 1:1}.
  • Number -1: prefix 0, need 0, found once, count 1. Dict becomes {0:2, 1:1}.
  • Number 2: prefix 2, need 2, not in the dict, count stays 1. Dict becomes {0:2, 1:1, 2:1}.
  • Number -2: prefix 0, need 0, found twice, count 3.

The answer is 3. The subarrays are [1, -1], [2, -2] and [1, -1, 2, -2]. Check it with count_slow before you trust the dict version.

What you now know

  • The pattern is to ask the dict a question first, then record the item, so an item never pairs with itself.
  • Two-sum runs in O(n) with a dict and in O(n log n) with sorting and two pointers, and the choice depends on whether the list is sorted and whether you need indexes.
  • Frequency counting uses Counter, and a dict keeps first-seen order when you need it.
  • Anagram groups use a sorted-letters key, and a list of 26 integers checks a single pair.
  • The longest run uses a set and starts counting only from values whose lower neighbour is missing, which keeps the work linear.
  • Subarray sums with negative numbers use prefix sums, a dict of earlier prefix counts, and a starting entry of 0 mapped to 1.
  • Python integers do not overflow, so counts and sums need no special type, but a list of many numbers still costs memory for each object.
Practice this topic: coding problems for Dicts and Sets →

Stuck? Ask the tutor

Ask about anything in this lesson. Answers are short and use simple words.

Sign in to ask the tutor.

← Dicts, sets and sorted keys in Python