Counting and grouping

Many problems become one pass over the data once you have a dict or a set. This lesson collects the patterns that come up again and again: counting, grouping, looking for a partner, and remembering prefix sums.

Counting

To count how often each item appears, keep a dict from item to count. Step through the first variant of the figure. The second variant comes up in the next sections.

Python offers three ways to write the same loop:

counts = {}
for w in words:
    counts[w] = counts.get(w, 0) + 1     # get returns 0 for a new key

from collections import defaultdict
counts = defaultdict(int)                # a missing key starts at int(), 0
for w in words:
    counts[w] += 1

from collections import Counter
counts = Counter(words)                  # the same, in one line
counts.most_common(2)                    # [('to', 2), ('be', 2)]

All three take O(n) time on average for n items. most_common(k) picks the k largest counts with a heap (HEAP station), which costs O(d log k) for d different items, cheaper than sorting all of them.

Grouping by a key

To put items into groups, choose a group key: a hashable value that is the same for items in the same group and different for items in different groups. Then keep a dict from group key to a list of items.

Example: group words that are anagrams of each other (the same letters in a different order, like listen and silent). Two words are anagrams exactly when their sorted letters are the same. So the sorted letters make a good group key:

from collections import defaultdict

def anagram_groups(words):
    groups = defaultdict(list)
    for w in words:
        groups[''.join(sorted(w))].append(w)    # 'listen' -> 'eilnst'
    return list(groups.values())

Sorting a word of length L costs O(L log L). A faster key counts the letters instead: a tuple of 26 counts, built in O(L). The key must be a tuple, not a list, because a dict key must be hashable (lesson 1).

Looking for a partner

Some problems ask for two items that fit together: two values that add up to a target, or two words that are mirror images. Trying every pair is O(n²). Instead, walk once, and keep the items already passed in a dict or set. For each new item, work out the exact partner it needs, and check whether that partner was already seen. Step through the second variant of the figure.

def pair_with_sum(a, target):
    seen = {}                        # value -> its place
    for i, x in enumerate(a):
        if target - x in seen:
            return (seen[target - x], i)
        seen[x] = i                  # store after the check
    return None

The check comes before the store. Otherwise, with target 12, a single 6 would pair with itself. The cost is one pass with O(1) dict work per item: O(n) on average, and O(n) extra space for seen.

Prefix sums in a dict

ARRS-4 used prefix sums: p[j] is the total of the first j values, and the stretch of values after place i up to place j adds up to p[j] − p[i]. Combine this with the partner idea. A stretch ending at place j adds up to k exactly when some earlier prefix sum equals p[j] − k. That earlier prefix sum is the partner.

So walk once, keeping a running total, and keep a dict that counts how many times each prefix sum has appeared. At each place, the number of good stretches ending there is a single dict lookup. Practice HASH-P4 asks you to finish this idea. It works even with negative values, where the sliding window from ARRS-3 fails, because a window’s total no longer grows as the window grows.