Skip to content

Intermediate · Errors, modules and the standard library · topic 25 of 50 · 25 min

collections and itertools

Explanation

The collections module gives you smarter containers: Counter counts things, defaultdict fills in missing keys, namedtuple gives a tuple named fields, and deque is a fast queue. The itertools module has ready-made loop tools: chain, count, combinations, permutations, product and groupby. They replace many long hand-written loops.

Counter: counting made easy

python
from collections import Counter
votes = ["chai", "coffee", "chai", "lassi", "chai", "coffee"]
count = Counter(votes)
print(count["chai"])
print(count.most_common(2))
print(Counter("banana"))
Shows 3, then [('chai', 3), ('coffee', 2)], then the letter counts.
Think of it like this: A Counter is like the tally marks a class monitor makes on the blackboard during a class election: one mark per vote, and at the end you read who has the most.

defaultdict: no more KeyError

python
from collections import defaultdict
people = [("Asha", "Pune"), ("Ravi", "Delhi"), ("Meena", "Pune")]
by_city = defaultdict(list)
for name, city in people:
    by_city[city].append(name)   # no need to check if city exists
print(dict(by_city))

spend = defaultdict(int)
spend["food"] += 250
print(spend["food"], spend["travel"])
A missing key gets a fresh value: list() gives [], int() gives 0.

namedtuple: tuples with names

python
from collections import namedtuple
Student = namedtuple("Student", ["name", "city", "marks"])
s = Student("Asha", "Pune", 88)
print(s.name, s.marks)
print(s[1])
s.marks is much clearer than s[2]. It is still a tuple, so s[1] works too.

deque: a queue open at both ends

python
from collections import deque
recent = deque(maxlen=3)
for search in ["phone", "charger", "cover", "earphones"]:
    recent.append(search)
print(list(recent))

queue = deque(["Ravi", "Asha"])
queue.append("Meena")      # joins at the back
print(queue.popleft())     # served from the front
maxlen keeps only the newest items; popleft is fast, unlike list.pop(0).

itertools: tools for looping

FunctionGives youExample result
count(1)1, 2, 3, ... forevernumbering with zip
chain(a, b)a then b, as one sequencechain([1, 2], [3]) → 1, 2, 3
combinations(x, 2)groups where order does not matterAB, AC, BC
permutations(x, 2)arrangements where order mattersAB, AC, BA, BC, CA, CB
product(a, b)every pair, one from eachS-red, S-blue, M-red, M-blue
groupby(x, key)runs of neighbouring items with the same keyneeds sorted data
python
from itertools import count, chain, combinations, permutations, product
for n, name in zip(count(1), ["Asha", "Ravi", "Meena"]):
    print(n, name)
print(list(chain([1, 2], [3], [4, 5])))
print(list(combinations(["A", "B", "C"], 2)))
print(len(list(permutations(["A", "B", "C"], 2))))
print(list(product(["S", "M"], ["red", "blue"])))
Most itertools functions give items one by one; list() collects them so you can print them.
Tip: Pick a team of 2 from 4 players: order does not matter, so use combinations (6 ways). Give a first and a second prize to 2 of 4 players: order matters, so use permutations (12 ways).

groupby: sort first!

python
from itertools import groupby
sales = [("fruit", 40), ("veg", 30), ("fruit", 25), ("veg", 10)]
sales.sort(key=lambda s: s[0])
for category, rows in groupby(sales, key=lambda s: s[0]):
    total = 0
    for _, amount in rows:
        total += amount
    print(category, total)
After sorting: fruit 65, veg 40.
Careful: groupby only joins items that sit next to each other. Without sorting, ["fruit", "veg", "fruit"] makes three groups: fruit, veg and fruit again. Always sort by the same key before groupby. Also never call list(count()): count never ends.

Watch the lesson

collections and itertools · 1/6

collections · itertools

Python has two modules full of ready-made tools for counting, grouping and looping. They save you a lot of code.

Prefer a video by a teacher? Find free videos on “python collections counter defaultdict itertools tutorial” on YouTube ↗

In real life: Which items sell most at a kirana shop?

Farhan writes down every item sold in a day. At closing time he wants the top three items and the money taken per category. Counter answers the first question in one line, and defaultdict adds up the money without any "if key not in dict" checks.

python
from collections import Counter, defaultdict
sold = [("milk", "dairy", 30), ("bread", "bakery", 40), ("milk", "dairy", 30),
        ("curd", "dairy", 25), ("milk", "dairy", 30), ("bread", "bakery", 40)]
items = Counter()
for name, _, _ in sold:
    items[name] += 1
print(items.most_common(3))
money = defaultdict(int)
for _, category, price in sold:
    money[category] += price
print(dict(money))
Top items, then {'dairy': 115, 'bakery': 80}.

Practice (0/3 done)

Task 1 · Code · Python

Write top_items(items, n) that returns the n most common items with their counts, as a list of (item, count) pairs, most common first. Use Counter.

Task 2 · Code · Python

Write group_by_city(people) where people is a list of (name, city) pairs. Return a normal dict that maps each city to the list of names from that city, in the order they appear. Use defaultdict(list).

Task 3 · Code · Python

Write team_pairs(players) that returns a list of every possible pair of players as tuples, where order does not matter (A with B is the same team as B with A). Use itertools.

Quick quiz

Question 1

1.What is Counter("banana")["a"]?

Question 2

2.With d = defaultdict(int), what happens on d["x"] += 1 when "x" is not a key yet?

Question 3

3.How many items does list(combinations(["A", "B", "C", "D"], 2)) have?

Question 4

4.You run groupby on ["veg", "fruit", "veg"] without sorting. How many groups do you get?

Question 5

5.How many pairs does product(["S", "M"], ["red", "blue"]) give?

Read more (free)

Pass the quiz (60% or more) or mark this topic as done when you are ready.

Your progress is saved on this device only.

Next: JSON and regular expressions →

← Previous: Dates and times with datetime