Sliding window with a frequency dict in Python

Why this matters

Many programming questions ask about a part of a longer sequence. The sequence might be a line of text, a list of numbers or a list of requests. The part you want is called a window: a stretch that has a start and an end. Finding the best window by checking every start and every end is slow. A long text with n letters has about n times n divided by 2 stretches, and checking each one again costs more time.

A sliding window avoids that. It keeps two positions, a left end and a right end. The right end moves forward one step at a time and adds one letter to the window. The left end moves forward only when the window breaks a rule. Each letter is added once and removed at most once, so the total work stays close to the length of the text.

This lesson uses a dict as the summary of what is inside the window. A dict stores pairs of a key and a value and finds a value quickly when you give it the key. Here the key is a letter and the value is how many times that letter appears in the window. A second number, the distinct counter, tracks how many different letters are in the window right now.

The problem is this. Given a text s and a limit k, find the length of the longest stretch of s that has at most k different letters. For the text eceba and k = 2, the answer is 3, because ece has two different letters and no longer stretch does.

A small example first

Before any code, find the answer by hand. The text is eceba and the limit is 2.

• e has 1 different letter. Length 1.
• ec has 2 different letters. Length 2.
• ece has 2 different letters, e and c. Length 3.
• eceb has 3 different letters. Too many.
• ceb has 3 different letters. Too many.
• eb has 2 different letters. Length 2.
• ba has 2 different letters. Length 2.

The longest valid stretch is ece with length 3.

Checking every stretch by hand is easy for five letters. For a text with a million letters it is not. The rest of the lesson shows how the window does the same job in one pass.

Python strings support len(s), and s[3] returns the character at index 3. Indexes start at 0, so the first character is at index 0.

The window idea

The window is the letters from index left up to and including index right. Its length is right - left + 1. For example, if left is 2 and right is 4, the window holds three letters: the letters at indexes 2, 3 and 4.

Each step of the loop does four things:

1. Add the letter at right to the counts. If it was not in the window before, increase the distinct counter.
2. If the window now has more than k distinct letters, remove the letter at left from the counts. If its count drops to 0, delete the key from the dict and decrease the distinct counter. Move left one step forward. Repeat until the window is valid again.
3. Compare the window length with the best length found so far.
4. Move on to the next right.

The window never moves backward. right goes from 0 to the last index, and left only goes forward. So the loop does about twice as many steps as there are letters, not one step for every possible stretch.

The moves can be counted. Running the window over eceba moves the right end 5 times, once for each letter, and the left end 3 times. That is 8 moves for a 5-letter text. Each letter is added once and removed at most once, so the number of moves never goes above twice the length of the text.

Step through eceba with k = 2

The counts are written the way Python prints a dict, in backticks so they are easy to read. Each line shows the state after the step.

• Index 0, letter e. Counts are {'e': 1}. Distinct is 1. Window is e. Best is 1.
• Index 1, letter c. Counts are {'e': 1, 'c': 1}. Distinct is 2. Window is ec. Best is 2.
• Index 2, letter e. The letter was already in the window, so only its count goes up. Counts are {'e': 2, 'c': 1}. Distinct is still 2. Window is ece. Best is 3.
• Index 3, letter b. Counts are {'e': 2, 'c': 1, 'b': 1}. Distinct is 3, which is too many. Remove the letter at left (index 0, e). Its count goes from 2 to 1, so the key stays. Distinct is still 3. Move left to 1. Remove the letter at left (index 1, c). Its count goes from 1 to 0, so the key is deleted and distinct drops to 2. Move left to 2. Window is eb. Best stays 3.
• Index 4, letter a. Counts are {'e': 1, 'b': 1, 'a': 1}. Distinct is 3, which is too many. Remove the letter at left (index 2, e). Its count goes from 1 to 0, so the key is deleted and distinct drops to 2. Move left to 3. Window is ba. Best stays 3.

The loop ends with best equal to 3.

The Python code

The code below is the whole program. The longest_at_most_k function is the window. The brute_force function checks every stretch the slow way. The test code runs five fixed examples and then compares the two functions on 2000 random strings. The random check is what gives confidence that the window code is correct, not only that it passes the examples.

eceba k=2 -> 3
aa k=1 -> 2
abcadcacacaca k=3 -> 11
abc k=0 -> 0
empty k=2 -> 0
random checks passed: 2000

The code is built from these parts.

counts = {} is the frequency dict. Its keys are letters, which are strings of length 1, and its values are counts. A dict can hold any hashable key, so the letters work directly with no wrapper type. Python has no separate boxing step, which is different from Java.

counts.get(c, 0) returns the count for c, or 0 if c is not in the dict. Using it avoids the KeyError described in the mistakes section below.

if before == 0 is the test for a new distinct letter. The letter was not in the window before this step, so the number of distinct letters goes up by one.

The while loop shrinks the window. It keeps going while distinct > k. A while is needed here because one removed letter may not be enough to bring the distinct count back down.

best = max(best, right - left + 1) records the window length after the window is valid.

The loop uses enumerate(s), which gives the index and the letter together. That removes the need for a separate counter.

Second example: more letters, a longer window

Take the text abcadcacacaca and k = 3. The answer is 11.

The window starts at index 0. It grows through a, b, c and a. After index 3 it holds abca, which has 3 distinct letters. That is still allowed.

At index 4 the letter is d. The window is abcad, with 4 distinct letters. The loop removes from the left. It removes a at index 0, but the count of a is still 1 after that, so the window bcad still has 4 letters. It removes b at index 1 and the window becomes cad, which has 3 letters. The left end is now at index 2.

From index 2 onward, the window only holds letters c, a and d, so it can grow to the end of the text. The window from index 2 through index 12 is cadcacacaca. It has 11 letters and only 3 distinct ones. No longer stretch works, so the answer is 11.

Third example: a shrink that takes two steps

Take the text aabbcc and k = 2. The answer is 4.

• Index 0 to 3, letters a, a, b, b. The window is aabb, with 2 distinct letters. Best is 4.
• Index 4, letter c. The window is aabbc, with 3 distinct letters. Too many. Remove the letter at index 0, which is a. Its count goes from 2 to 1, so the key stays. Distinct is still 3. Move left to 1. Remove the letter at index 1, which is also a. Its count goes from 1 to 0, so the key is deleted. Distinct drops to 2. Move left to 2. Window is bbc, with length 3. Best stays 4.
• Index 5, letter c. The window is bbcc, with 2 distinct letters and length 4. Best stays 4.

Here the right end added one letter but the left end had to move twice. A single if would not be enough. The while loop handles any number of removals.

Returning the stretch, not only its length

The function above returns a length. To return the stretch itself, store the start of the best window when you find it, then slice the text at the end. The > test, rather than >=, keeps the first window that reaches the best length.

aabb
aaa

The slice s[best_start:best_start + best] copies the window's letters into a new string. The stretch is aabb for the first example, which is the first of the two stretches with length 4.

Why the distinct counter exists

There is a shortcut that looks simpler. A dict has a len() that returns the number of keys. If the code deletes a key as soon as its count reaches 0, then len(counts) is the number of distinct letters in the window. The code could use that and skip the separate counter.

The counter is still useful for two reasons. First, it makes the rule easy to read: a new letter raises the count and a letter whose count reaches 0 lowers it. Second, it makes the code correct even if a later change keeps zero counts in the dict. A dict that keeps zero counts has a larger len() than the number of letters in the window, and the answer would be wrong. The mistakes section below shows what happens.

The counter must change in exactly two places: when a letter first enters the window, and when a letter's count drops to 0 on the left. If you change it anywhere else, the number goes out of step with the dict.

Common mistakes

A missing key raises KeyError

Reading a key that is not in a dict with square brackets raises KeyError. This is the Python version of the null problem in Java: the program tries to use a value that is not there.

counts = {"a": 1}
print(counts["z"])

KeyError: 'z'

The fix is counts.get("z", 0), as used in the code above. The get method returns the default when the key is missing, so the window code never needs to check first.

Keeping zero counts in the dict

If the shrink step lowers a count to 0 but leaves the key in the dict, len(counts) keeps counting that letter. The window then looks like it has more distinct letters than it does. A version of the code that uses len(counts) and does not delete keys fails on eceba:

def zero_kept(s, k):
    counts = {}
    left = 0
    best = 0
    for right, c in enumerate(s):
        counts[c] = counts.get(c, 0) + 1
        while len(counts) > k:
            out = s[left]
            counts[out] -= 1
            left += 1
        best = max(best, right - left + 1)
    return best
 
 
print(zero_kept("eceba", 2))

Traceback (most recent call last):
  ...
KeyError: 'a'

The dict kept the letter e with count 0 after it left the window, so the loop kept moving left forward. It reached letters such as a that had never been added to the dict, and counts[out] raised KeyError for them. The fix is to delete the key when its count reaches 0, as the code above does.

Shrinking with if instead of while

An if where the code needs a while can still print the right length for the examples in this lesson, so the length alone does not reveal the bug. The window itself can be left with too many distinct letters. The next program uses if and reports the final window and its distinct count with k = 1 on aabac.

('ac', 2)

The final window is ac, which has 2 distinct letters, but the limit is 1. The if removed one letter per step, and the window was not valid at the end. Use while so the window is valid at the end of every step.

Forgetting the length formula

The length of a window from left to right is right - left + 1, not right - left. The + 1 is there because both ends are included. Leaving it out makes each answer one shorter.

Practice

Exercise 1

Find the longest stretch of aaabbc that has at most 1 distinct letter. Give its length and the stretch.

Show the answer

The longest stretch is aaa, with length 3. The stretch bb has length 2 and bbc has 2 distinct letters, so neither is longer.

Exercise 2

Trace the window for abaccc with k = 2. At which index does the window first become invalid, and what is the final answer?

Show the answer

The window becomes invalid at index 3, the letter c. Before that the window is aba, with letters a and b. Adding c gives abac, with 3 distinct letters. The loop removes the a at index 0 (its count goes from 2 to 1), which still leaves 3 distinct letters, so it removes the b at index 1. The window becomes ac, which has 2 distinct letters. The final answer is 4, from the stretch accc. The program in this lesson prints 4 for this input.

Exercise 3

Change longest_at_most_k so it returns the best stretch as a string instead of its length. Keep the same counts and distinct logic.

Show the answer

The code is the longest_substring function in the section above. Store the best start whenever a longer valid window is found, then return the slice. For aabbcc with k = 2 the function returns aabb.

Exercise 4

Why does distinct stay the same when a letter that is already in the window is added?

Show the answer

The count for that letter was already 1 or more, so before is not 0 and the if before == 0 test fails. The letter is already counted as one distinct letter, and only its count goes up. For example, adding a second e to the window ec changes the counts from {'e': 1, 'c': 1} to {'e': 2, 'c': 1}, and distinct stays 2.

Exercise 5

What does longest_at_most_k return for an empty text with k = 2, and for abc with k = 0?

Show the answer

Both return 0. For the empty text the for loop never runs, so best stays 0. For abc with k = 0, every letter adds a new distinct letter, the while loop removes it again before the length is recorded, so the window is always empty of valid letters and best stays 0. The program prints empty k=2 -> 0 and abc k=0 -> 0.

Summary

• A sliding window keeps a left end and a right end. Both only move forward, so the work is about one step per letter.
• A dict from letter to count tells you what is in the window, and a separate distinct counter tells you how many different letters there are.
• The distinct counter goes up when a letter's count goes from 0 to 1 and down when it goes from 1 to 0.
• Shrink the window with a while loop, because one removal may not be enough to restore the rule.
• Delete a key from the dict when its count reaches 0, and use get with a default so a missing key does not raise KeyError.
• Check the window function against a slow brute-force version on random input before you trust it.