Data structures and idiomatic Python
¶

Alfred Galichon (NYU)
¶

'math+econ+code' masterclass series: fundamentals of research in python
¶

With python code examples
¶

© 2018–2026 by Alfred Galichon. Past and present support from NSF grant DMS-1716489, ERC grant CoG-866274 are acknowledged, as well as inputs from contributors listed here.

If you reuse material from this masterclass, please cite as:
Alfred Galichon, 'math+econ+code' masterclass series. https://www.math-econ-code.org/

Learning objectives¶

  • Distinguish the four built-in containers, list, tuple, dict, set, by mutability, ordering, and access pattern, and choose the right one for a given task.

  • Use comprehensions (list, dict, set, generator) to express a transformation of a sequence in one readable line.

  • Recognize Python's reference semantics: assignment binds names to objects and never copies, so mutating an object is visible through every name bound to it. This is the source of most beginner bugs, and of one exercise below.

  • Apply the everyday idioms, enumerate, zip, tuple and starred unpacking, dict merging, that separate idiomatic Python from Python transliterated out of C++.

  • Build an index map: a dictionary from a pair $(x,y)$ to a position in a flat vector. This is the bookkeeping that turns a matrix of economic primitives into the argument of a solver, and it is the explicit form of the row-major flattening $\operatorname{vec}_C$ that fd08 will hand to NumPy.

  • Solve a small assignment problem by exhaustive enumeration, and discover why a greedy rule gets it wrong.

References¶

[M] McKinney, W. (2022). Python for Data Analysis (3rd ed.). O'Reilly, Chapter 3. Online at https://wesmckinney.com/book/.

[V] VanderPlas, J. (2023). Python Data Science Handbook (2nd ed.). O'Reilly.

[QE] Sargent, T. J. and Stachurski, J. QuantEcon Python Lectures. https://python-programming.quantecon.org/.

[PEP8] van Rossum, G., Warsaw, B., and Coghlan, A. (2001). PEP 8: Style Guide for Python Code. https://peps.python.org/pep-0008/.

[G] Galichon, A. (2016). Optimal Transport Methods in Economics. Princeton University Press, Chapter 3: for the assignment problem of §12 and its dual.

[K] Koopmans, T. C. and Beckmann, M. (1957). "Assignment Problems and the Location of Economic Activities." Econometrica 25(1), 53–76: the assignment problem as an economic model, with prices.

1. Motivation¶

fd01 gave us numbers, names, and control flow: enough to write a calculation, but not enough to organize one. Almost every economic object we will compute with is indexed: a cash flow by date, a surplus by a pair (worker, firm), a transition probability by a triple (state, action, next state). Getting from those indices to something a solver will accept is, in practice, a large fraction of the work, and it is where quiet errors live.

Python's four built-in containers are the tools for that job, and this lecture is about choosing among them well. The last section puts them to work: we build an index map for a small assignment problem, use it to flatten a matrix of surpluses into a vector, and solve the problem by brute force. That index map is the same object fd08 will construct implicitly with reshape(-1), and the same one the ot and lp series use to assemble sparse constraint matrices. Meeting it here, written out by hand in a dictionary, is the point of doing this lecture before the libraries arrive.

2. Four built-in containers¶

Almost every Python program you read uses the same four containers, and choosing among them well is most of writing readable code.

Container Mutable? Ordered? Indexed by Typical use
list yes yes integer position a sequence you will modify
tuple no yes integer position a fixed record; a composite key; multiple return
dict yes yes (insertion order, since 3.7) hashable key a mapping or lookup table
set yes no : uniqueness, membership, set algebra

We take them one by one, then turn to the idioms, comprehensions, unpacking, zip, enumerate, that bind them together.

3. Lists¶

A list is a mutable, ordered sequence. It is the right default when you need some container and do not yet know which.

In [1]:
# nominal GDP (USD trillions), largest economies, 2022 (rounded)
countries = ["USA", "China", "Japan", "Germany", "UK", "France"]
gdp = [25.46, 17.96, 4.23, 4.07, 3.07, 2.78]

len(gdp), gdp[0], gdp[-1]
Out[1]:
(6, 25.46, 2.78)

Slicing. xs[a:b] returns a new list with the elements at positions $a, a+1, \dots, b-1$. The upper bound is exclusive, the same convention as range in fd01 §9. xs[a:b:s] steps by $s$, and a negative step reverses.

In [2]:
countries[:3], countries[1:3], countries[-2:], countries[::-1]
Out[2]:
(['USA', 'China', 'Japan'],
 ['China', 'Japan'],
 ['UK', 'France'],
 ['France', 'UK', 'Germany', 'Japan', 'China', 'USA'])

Mutation. Lists are mutable: they can be changed in place.

In [3]:
gdp.append(2.01)            # add Italy at the end
countries.append("Italy")
gdp[0] = 25.50              # revise the USA figure
gdp
Out[3]:
[25.5, 17.96, 4.23, 4.07, 3.07, 2.78, 2.01]

The standard mutators are append, extend, insert, pop, remove, and the in-place .sort(). Note the difference between the in-place method and the pure function:

  • xs.sort() mutates xs and returns None;
  • sorted(xs) returns a new sorted list and leaves xs alone.

This pattern, a mutating method on the object, versus a pure function returning a new one, runs through the whole language, and forgetting which is which is a reliable source of confusion. The cell below makes the distinction visible.

In [4]:
gdp_sorted = sorted(gdp, reverse=True)      # pure: returns a new list
print("sorted(gdp, reverse=True) ->", gdp_sorted)
print("gdp is unchanged          ->", gdp)

scratch = gdp.copy()                        # in-place: mutates, returns None
returned = scratch.sort()
print("\nscratch.sort() returned   ->", returned)
print("scratch was reordered     ->", scratch)
sorted(gdp, reverse=True) -> [25.5, 17.96, 4.23, 4.07, 3.07, 2.78, 2.01]
gdp is unchanged          -> [25.5, 17.96, 4.23, 4.07, 3.07, 2.78, 2.01]

scratch.sort() returned   -> None
scratch was reordered     -> [2.01, 2.78, 3.07, 4.07, 4.23, 17.96, 25.5]

Note that the in-place sort was applied to a copy, and deliberately so. countries and gdp are parallel lists: the correspondence between them is carried entirely by position, and nothing in the data enforces it. Sorting one of the two would silently destroy that correspondence: every subsequent line would still run, and every number it printed would be attached to the wrong country.

This is the central weakness of parallel lists, and the reason the next section reaches for a dictionary instead: a dict binds the country to its value in the data structure itself, so there is no correspondence left to break. In fd06 the same argument returns one level up, as the case for a DataFrame over a collection of parallel arrays.

Two lists concatenate with +, and a list repeats with *. Keep the second one in mind: it returns in Exercise 5, where it does something you may not expect.

In [5]:
["a", "b"] + ["c"], [0.0] * 4
Out[5]:
(['a', 'b', 'c'], [0.0, 0.0, 0.0, 0.0])

4. Tuples¶

A tuple is the immutable cousin of list. Use one when the collection has a fixed size or a fixed meaning: coordinates, a record, the multiple return value of a function, or, as in §12, a composite index $(x,y)$.

In [6]:
point = 3.14, 2.71
x, y = point                # unpacking
print(x, y)

nested = (1, (3, 4))        # unpacking follows the structure
a, (b, c) = nested
print(a, b, c)
3.14 2.71
1 3 4

The comma is what makes a tuple, not the parentheses:

In [7]:
singleton = (1,)        # a tuple of length 1
not_a_tuple = (1)       # just the integer 1, in parentheses
type(singleton), type(not_a_tuple)
Out[7]:
(tuple, int)

Multiple return. A function returns a tuple and the caller unpacks it: one of the most common Python idioms.

In [8]:
def summary(xs):
    return min(xs), max(xs), sum(xs) / len(xs)

lo, hi, avg = summary(gdp)
print(f"min = {lo:.2f},  max = {hi:.2f},  mean = {avg:.2f}")
min = 2.01,  max = 25.50,  mean = 8.52

Tuples are hashable (provided their entries are), so they can serve as dictionary keys. Lists cannot, precisely because they are mutable: a key that could change value underneath the dictionary would break the lookup. This is why a pair like (country, year), or the pair $(x,y)$ of a worker type and a firm type, is written as a tuple. It is the fact that §12 is built on.

In [9]:
panel = {("USA", 2022): 25.46, ("USA", 2021): 23.6, ("Japan", 2022): 4.23}
print(panel[("USA", 2022)])

try:
    {["USA", 2022]: 25.46}          # a list key: not hashable
except TypeError as err:
    print("TypeError:", err)
25.46
TypeError: unhashable type: 'list'

5. Dictionaries¶

A dict maps hashable keys to arbitrary values, with $O(1)$ average-case lookup, insertion, and deletion. It is the workhorse of Python.

In [10]:
# entered in alphabetical order -- note that this is NOT the order of magnitude
gdp_2022 = {
    "China": 17.96, "France": 2.78, "Germany": 4.07,
    "Japan": 4.23, "UK": 3.07, "USA": 25.46,
}
gdp_2022["USA"], len(gdp_2022), list(gdp_2022)
Out[10]:
(25.46, 6, ['China', 'France', 'Germany', 'Japan', 'UK', 'USA'])

Iterating a dict yields its keys; .values() yields the values and .items() yields (key, value) pairs. list(d) is therefore the list of keys.

Lookup with a default. d[k] raises KeyError if k is absent; d.get(k, default) returns default instead. Use .get when a missing key is a normal case rather than a bug, and let the KeyError fire when it is genuinely an error: a crash at the point of the mistake is worth more than a silent zero that propagates into your results.

In [11]:
print("present, [] :", gdp_2022["USA"])
print("absent,  .get :", gdp_2022.get("Brazil", 0.0))
print("membership    :", "USA" in gdp_2022, "|", "Brazil" in gdp_2022)

try:
    gdp_2022["Brazil"]
except KeyError as err:
    print("absent,  []   : KeyError", err)
present, [] : 25.46
absent,  .get : 0.0
membership    : True | False
absent,  []   : KeyError 'Brazil'

Iteration, and insertion order. Since Python 3.7 a dict preserves the order in which keys were inserted. That is convenient, but do not mistake it for being sorted: sort explicitly when you need an order with meaning.

In [12]:
for country, value in gdp_2022.items():
    print(f"{country:<8}  ${value:>5.2f}T")

print("\nsorted by GDP, descending:")
for country in sorted(gdp_2022, key=gdp_2022.get, reverse=True):
    print(f"{country:<8}  ${gdp_2022[country]:>5.2f}T")
China     $17.96T
France    $ 2.78T
Germany   $ 4.07T
Japan     $ 4.23T
UK        $ 3.07T
USA       $25.46T

sorted by GDP, descending:
USA       $25.46T
China     $17.96T
Japan     $ 4.23T
Germany   $ 4.07T
UK        $ 3.07T
France    $ 2.78T

The key=gdp_2022.get argument deserves a second look: sorted is being handed the dictionary's own lookup method as the function that ranks the keys. Passing a function as an argument to another function is routine in Python, and fd03 treats it properly.

6. Sets¶

A set is an unordered collection of unique, hashable elements. Two reasons to use one: deduplication, and set algebra. In §12 a set is what holds the support of a matching: the collection of pairs that are actually matched.

In [13]:
g7 = {"USA", "Canada", "UK", "France", "Germany", "Italy", "Japan"}
g20 = g7 | {"China", "India", "Brazil", "Russia", "Mexico", "Indonesia",
            "Saudi Arabia", "Argentina", "Turkey", "South Africa",
            "South Korea", "Australia"}

print("sizes                :", len(g7), len(g20))
print("G7 is a subset of G20:", g7 <= g20)
print("intersection = G7    :", (g7 & g20) == g7)
print("G20 but not G7       :", sorted(g20 - g7))
print("in exactly one       :", len(g7 ^ g20))
sizes                : 7 19
G7 is a subset of G20: True
intersection = G7    : True
G20 but not G7       : ['Argentina', 'Australia', 'Brazil', 'China', 'India', 'Indonesia', 'Mexico', 'Russia', 'Saudi Arabia', 'South Africa', 'South Korea', 'Turkey']
in exactly one       : 12

Duplicates simply collapse: a set has no notion of multiplicity:

In [14]:
{"USA", "France", "Japan", "France"}, set([1, 2, 4, 2])
Out[14]:
({'France', 'Japan', 'USA'}, {1, 2, 4})

Membership testing is $O(1)$ for sets and dicts, but $O(n)$ for lists. For a collection of any size this is the difference between a fast loop and a slow one, and it is the most common performance mistake in otherwise correct research code: an inner loop that repeatedly asks if x in some_list. We measure the difference in fd04; for now, take it as a rule: if you are testing membership, use a set.

In [15]:
# deduplicate a list while preserving first-appearance order
items = ["b", "a", "c", "a", "b", "d"]

seen = set()
unique = [x for x in items if not (x in seen or seen.add(x))]
print("order-preserving unique:", unique)

# clearer, and does the same thing: dict keys are unique and ordered
print("via dict.fromkeys     :", list(dict.fromkeys(items)))
order-preserving unique: ['b', 'a', 'c', 'd']
via dict.fromkeys     : ['b', 'a', 'c', 'd']

The first version works because seen.add(x) always returns None, which is falsy, so the or short-circuits and the element is kept exactly when it was not already seen. It is a cute idiom and you will meet it in the wild, but dict.fromkeys says the same thing without the puzzle, and, per PEP 20, readability counts.

7. Mutability and references¶

Every variable in Python is a name bound to an object. Assignment binds; it never copies. This is the single most important fact about the language, and the source of most beginner bugs, including one that has cost researchers published corrections.

In [16]:
xs = [1, 2, 3]
ys = xs                 # ys and xs name the SAME list
ys.append(99)

print("xs        =", xs)
print("ys        =", ys)
print("same object?", xs is ys, "| ids:", id(xs) == id(ys))
xs        = [1, 2, 3, 99]
ys        = [1, 2, 3, 99]
same object? True | ids: True

is asks whether two names refer to the same object; == asks whether two objects have equal value. They are different questions, and equal value does not imply identity:

In [17]:
list1 = [1, 2, 5]
list2 = [1, 2, 5]
print("equal in value? ", list1 == list2)
print("same object?    ", list1 is list2)
equal in value?  True
same object?     False

To get an actual copy, ask for one: list(xs), xs[:], and xs.copy() all produce a new top-level list. (For nested structures these are shallow copies, the inner objects are still shared, and copy.deepcopy is the tool when that matters.)

In [18]:
xs = [1, 2, 3]
ys = xs.copy()          # or list(xs), or xs[:]
ys.append(99)

print("xs =", xs, " unchanged")
print("ys =", ys)
xs = [1, 2, 3]  unchanged
ys = [1, 2, 3, 99]

Immutable objects, int, float, str, tuple, frozenset, sidestep the issue entirely. There is no way to mutate the integer 5, so sharing it is harmless. The trap is the mutables: list, dict, set, and most objects you define yourself.

Mutable default arguments: the classic gotcha. A function's default arguments are evaluated once, when the def statement runs. If a default is a mutable object, every call that omits that argument shares the same object, and state leaks across calls that were meant to be independent.

In [19]:
def buggy_log(message, history=[]):        # do not do this
    history.append(message)
    return history

print(buggy_log("first"))
print(buggy_log("second"))
print(buggy_log("third"))                  # the list keeps growing across calls
['first']
['first', 'second']
['first', 'second', 'third']

The fix is to use None as a sentinel and build a fresh container inside the function:

In [20]:
def good_log(message, history=None):
    if history is None:
        history = []
    history.append(message)
    return history

print(good_log("first"))
print(good_log("second"))                  # a fresh list on each call
['first']
['second']

It is worth being precise about why this matters beyond tidiness. In a simulation or an estimation loop, an accidentally shared mutable is a channel through which one iteration contaminates the next. The result is not a crash but a number: plausible, reproducible, and wrong. Nothing in the output announces it. This is exactly the class of error that the verification habit of fd01 is meant to catch, and that the tests of fd04 are meant to catch automatically.

8. Comprehensions¶

A list comprehension is a for loop that builds a list inline. It is the most distinctive piece of Python syntax, and adopting it makes code noticeably shorter and, past a short adjustment period, easier to read.

In [21]:
# the verbose form
squares = []
for k in range(10):
    squares.append(k ** 2)

# the same thing, as a comprehension
squares_comp = [k ** 2 for k in range(10)]

print(squares_comp)
print("agree:", squares == squares_comp)
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
agree: True

A filter clause keeps only some elements:

In [22]:
[k ** 2 for k in range(10) if k % 2 == 0]
Out[22]:
[0, 4, 16, 36, 64]

An economic example. Year-on-year growth rates from a level series. Note the index arithmetic: the growth series has one fewer entry than the level series, which is the discrete-time counterpart of differentiation losing a boundary point.

In [23]:
gdp_series = [21.4, 21.5, 23.6, 25.7, 27.7]      # USA, 2019-2023, USD trillions

growth = [(gdp_series[t] - gdp_series[t - 1]) / gdp_series[t - 1]
          for t in range(1, len(gdp_series))]

print("levels :", gdp_series, f"({len(gdp_series)} entries)")
print("growth :", [f"{g:+.1%}" for g in growth], f"({len(growth)} entries)")
levels : [21.4, 21.5, 23.6, 25.7, 27.7] (5 entries)
growth : ['+0.5%', '+9.8%', '+8.9%', '+7.8%'] (4 entries)

Dict and set comprehensions use the same idea with different brackets:

In [24]:
population_2022 = {"USA": 333, "China": 1412, "Japan": 125,
                   "Germany": 84, "UK": 67, "France": 68}       # millions

# dict comprehension: GDP per capita, in USD
gdp_per_capita = {c: gdp_2022[c] * 1e6 / population_2022[c] for c in gdp_2022}

for c, v in sorted(gdp_per_capita.items(), key=lambda kv: -kv[1]):
    print(f"{c:<8} ${v:>9,.0f}")

# set comprehension: the distinct initials
print("\ninitials:", {c[0] for c in gdp_2022})
USA      $   76,456
Germany  $   48,452
UK       $   45,821
France   $   40,882
Japan    $   33,840
China    $   12,720

initials: {'C', 'J', 'U', 'G', 'F'}

Read this ranking against the level ranking of §5: China is second by GDP and last by GDP per capita. The two orderings answer different questions, aggregate size versus average living standard, and a good deal of bad economic commentary consists of quoting one and drawing conclusions about the other.

Two rules of thumb. Comprehensions can be abused. If you find yourself nesting them more than two deep, or stacking several filters, write an ordinary for loop instead. Readability wins.

9. Generators and iterators¶

A generator expression uses parentheses where a list comprehension uses square brackets, and produces its values lazily: nothing is computed until something asks for it. The difference is not stylistic: it is the difference between holding a sequence in memory and not.

In [25]:
import sys

n = 1_000_000
as_list = [k * k for k in range(n)]        # materialized: every value exists at once
as_gen = (k * k for k in range(n))         # lazy: nothing computed yet

print(f"list      : {sys.getsizeof(as_list):>10,} bytes")
print(f"generator : {sys.getsizeof(as_gen):>10,} bytes")
print(f"ratio     : {sys.getsizeof(as_list) / sys.getsizeof(as_gen):>10,.0f}x")

print("\nsame answer:", sum(as_list) == sum(k * k for k in range(n)))
list      :  8,448,728 bytes
generator :        200 bytes
ratio     :     42,244x
same answer: True

The generator is some forty thousand times smaller because it stores a rule, not a sequence. For a one-shot consumption like sum, max, or any, it is strictly better: the same answer, none of the memory. A generator is also exhausted once consumed, sum(as_gen) a second time returns 0, which is the price of laziness and a classic source of confusion.

The yield keyword turns a function into a generator: it produces values one at a time and remembers its state in between.

In [26]:
def fibs(n):
    a, b = 0, 1
    for _ in range(n):
        yield a
        a, b = b, a + b

list(fibs(10))
Out[26]:
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

Reach for a generator when (i) the sequence is large and you will scan it once, (ii) the sequence is unbounded, or (iii) you want to chain transformations without materializing the intermediate results. Otherwise a list is usually clearer, and, unlike a generator, it can be inspected, re-read, and printed while you are debugging.

10. The everyday idioms: enumerate, zip, unpacking¶

Three small things that account for an outsized share of whether code reads as Python or as transliterated C++.

enumerate when you need the index and the value:

In [27]:
for i, country in enumerate(countries[:4]):
    print(f"{i}: {country}")

# enumerate can start from 1, which is often what the mathematics wants
print("\n", list(enumerate(countries[:3], start=1)))
0: USA
1: China
2: Japan
3: Germany

 [(1, 'USA'), (2, 'China'), (3, 'Japan')]

Avoid for i in range(len(countries)) followed by countries[i]. It works, but it is noisier and it is the construction that produces off-by-one errors.

zip when you need to walk two or more sequences in lockstep. Note that zip stops at the shortest argument, silently, which is convenient when you mean it and a trap when you do not.

In [28]:
populations = [333, 1412, 125, 84, 67, 68]

for country, g, p in zip(countries, gdp, populations):
    print(f"{country:<8} GDP=${g:>5.2f}T  pop={p:>5}M  per capita=${g * 1e6 / p:>7,.0f}")

print("\nzip truncates to the shortest:", list(zip([2], [3, 4])))
USA      GDP=$25.50T  pop=  333M  per capita=$ 76,577
China    GDP=$17.96T  pop= 1412M  per capita=$ 12,720
Japan    GDP=$ 4.23T  pop=  125M  per capita=$ 33,840
Germany  GDP=$ 4.07T  pop=   84M  per capita=$ 48,452
UK       GDP=$ 3.07T  pop=   67M  per capita=$ 45,821
France   GDP=$ 2.78T  pop=   68M  per capita=$ 40,882

zip truncates to the shortest: [(2, 3)]

Note what the truncation just did: countries and gdp have seven entries after §3 appended Italy, but populations has six, so Italy was dropped without a word. If you want the mismatch to be an error rather than a silent omission, pass strict=True (Python 3.10+).

In [29]:
print("lengths:", len(countries), len(gdp), len(populations))

try:
    list(zip(countries, gdp, populations, strict=True))
except ValueError as err:
    print("with strict=True ->  ValueError:", err)
lengths: 7 7 6
with strict=True ->  ValueError: zip() argument 3 is shorter than arguments 1-2

Unpacking generalizes from a, b = pair to any sequence; a starred name absorbs the rest.

In [30]:
first, *middle, last = [1, 2, 3, 4, 5]
print(first, middle, last)
1 [2, 3, 4] 5

Dict merging with ** unpacking: the standard way to layer a set of overrides on top of a set of defaults, which is how most numerical routines are configured:

In [31]:
defaults = {"r": 0.05, "T": 30, "tol": 1e-8}
overrides = {"T": 50, "verbose": True}

config = {**defaults, **overrides}          # later keys win
config
Out[31]:
{'r': 0.05, 'T': 50, 'tol': 1e-08, 'verbose': True}

11. The * and ** operators, side by side¶

Both * and ** are unpacking operators, and the whole difference between them is what they operate on:

  • * spreads out the elements of an iterable: positional items, in order;
  • ``** spreads out the key–value pairs of a mapping.
operator operates on spreads out typical uses
* any iterable (list, tuple, range, …) its items, in order [*a, *b], first, *rest = seq, print(*words)
** a dict (mapping) its key: value pairs {**d1, **d2}

So * is for positional things and ** is for named things. That is exactly why a list cannot be **-unpacked, it has no keys, and why two dicts cannot be merged with *.

In [32]:
a, b = [1, 2], [3, 4]
print([*a, *b])              # * spreads list ELEMENTS      -> [1, 2, 3, 4]

words = ["alpha", "beta", "gamma"]
print(*words)                # * spreads into separate args -> alpha beta gamma
print(words)                 #   contrast: the list as ONE object

d1, d2 = {"r": 0.05, "T": 30}, {"T": 50, "tol": 1e-8}
print({**d1, **d2})          # ** spreads key-value PAIRS   -> later keys win
[1, 2, 3, 4]
alpha beta gamma
['alpha', 'beta', 'gamma']
{'r': 0.05, 'T': 50, 'tol': 1e-08}

The same two operators reappear in function definitions and calls: def f(*args, **kwargs) collects extra positional and keyword arguments into a tuple and a dict, while f(*seq, **dct) passes a sequence and a dict as arguments. That use is the subject of fd03.

12. Worked example: an index map for a matching problem¶

Everything so far now earns its keep. Consider the simplest interesting economic problem with a two-dimensional index: an assignment problem. There are worker types $x \in \mathcal{X}$ and firm types $y \in \mathcal{Y}$, and a match between $x$ and $y$ generates a surplus $\Phi_{xy}$. With one worker of each type and one position at each firm, the planner chooses which worker goes where so as to maximize total surplus.

This is the discrete optimal transport problem of ot01, and the object $\Phi_{xy}$ is the one the whole ot and et series is built around. Here we are not going to solve it properly, that needs the linear programming of lp01, but we will build the data structures it is stated in, which is the part that belongs to this lecture.

Three container choices, each for a reason established above:

  • the surplus is a dict keyed by tuples, Phi_x_y[(x, y)], because the index is a pair and tuples are hashable (§4);
  • the index map is a dict from a pair to an integer position, which is what lets a two-dimensional object be handed to a solver that only accepts vectors;
  • a candidate matching is a set of pairs (§6), because what matters is which pairs are in it, without order or multiplicity.
In [33]:
# Worker types and firm types. Following the house convention, an object carrying
# axes (x, y) is named with those axes, and nbx, nby count them.
X = ["analyst", "engineer", "manager"]
Y = ["bank", "startup", "university", "government"]
nbx, nby = len(X), len(Y)

# Surplus of a match, in arbitrary units of output
Phi_x_y = {
    ("analyst",  "bank"): 10.0, ("analyst",  "startup"): 9.0,
    ("analyst",  "university"): 5.0, ("analyst",  "government"): 4.0,
    ("engineer", "bank"): 9.0,  ("engineer", "startup"): 6.0,
    ("engineer", "university"): 4.0, ("engineer", "government"): 3.0,
    ("manager",  "bank"): 3.0,  ("manager",  "startup"): 4.0,
    ("manager",  "university"): 8.0, ("manager",  "government"): 7.0,
}

print(f"nbx = {nbx} worker types, nby = {nby} firm types, {len(Phi_x_y)} pairs\n")
print(f"{'':<10}" + "".join(f"{y:>12}" for y in Y))
for x in X:
    print(f"{x:<10}" + "".join(f"{Phi_x_y[(x, y)]:>12.1f}" for y in Y))
nbx = 3 worker types, nby = 4 firm types, 12 pairs

                  bank     startup  university  government
analyst           10.0         9.0         5.0         4.0
engineer           9.0         6.0         4.0         3.0
manager            3.0         4.0         8.0         7.0

The index map¶

A solver does not accept a dictionary keyed by pairs of strings; it accepts a vector. So we need a bijection between the pairs $(x,y)$ and the positions $0, 1, \dots, n_x n_y - 1$ of a flat vector. The convention this series uses throughout is row-major order: $y$ varies fastest, so the pair $(x_i, y_j)$ sits at position

$$ \operatorname{idx}(i,j) \;=\; i \cdot n_y + j . $$

This is exactly the flattening written $\operatorname{vec}_C$ in the charter, the $C$ standing for the row-major (C-language) ordering that NumPy's reshape uses by default. In fd08 it will be a single call, Phi_xy = Phi_x_y.reshape(-1); here we build it by hand, because doing it once by hand is what makes the single call auditable later.

In [34]:
# The index map: a pair -> its position in the flat vector
idx_x_y = {(x, y): i * nby + j for i, x in enumerate(X) for j, y in enumerate(Y)}

# The flat surplus vector, in that order. Per the naming convention, the
# concatenated suffix xy denotes the row-major flattening of the axes (x, y).
Phi_xy = [Phi_x_y[xy] for xy in sorted(idx_x_y, key=idx_x_y.get)]

print("index map:")
for xy, position in list(idx_x_y.items())[:5]:
    print(f"  {str(xy):<28} -> {position}")
print("  ...")

print(f"\nPhi_xy = {Phi_xy}")
print(f"length = {len(Phi_xy)} = nbx * nby = {nbx * nby}")
index map:
  ('analyst', 'bank')          -> 0
  ('analyst', 'startup')       -> 1
  ('analyst', 'university')    -> 2
  ('analyst', 'government')    -> 3
  ('engineer', 'bank')         -> 4
  ...

Phi_xy = [10.0, 9.0, 5.0, 4.0, 9.0, 6.0, 4.0, 3.0, 3.0, 4.0, 8.0, 7.0]
length = 12 = nbx * nby = 12

Verification. Two independent routes to the same ordering. The first walks the index map and sorts by position; the second simply iterates the two loops in the nested order $x$ outer, $y$ inner, which is what row-major means. If the index map is correct, they agree, and, just as importantly, the round trip from a position back to its pair must recover the original.

In [35]:
# route 1: order induced by the index map (above)
order_from_map = sorted(idx_x_y, key=idx_x_y.get)

# route 2: the nested loop, x outer and y inner -- the definition of row-major
order_from_loops = [(x, y) for x in X for y in Y]

print("orderings agree:", order_from_map == order_from_loops)
assert order_from_map == order_from_loops, "the index map is not row-major"

# route 3: the inverse map must undo the forward map, for every pair
inverse = {position: xy for xy, position in idx_x_y.items()}
round_trip_ok = all(inverse[idx_x_y[xy]] == xy for xy in idx_x_y)
print("round trip idx -> inverse -> idx recovers every pair:", round_trip_ok)
assert round_trip_ok

# and the flattening must agree entry by entry with the nested loop
gap = max(abs(Phi_xy[idx_x_y[(x, y)]] - Phi_x_y[(x, y)]) for x in X for y in Y)
print(f"max |flat[idx(x,y)] - Phi[x,y]| = {gap:.1e}   (tolerance 0, this is exact)")
assert gap == 0.0
print("check passed: the flattening is a faithful bijection.")
orderings agree: True
round trip idx -> inverse -> idx recovers every pair: True
max |flat[idx(x,y)] - Phi[x,y]| = 0.0e+00   (tolerance 0, this is exact)
check passed: the flattening is a faithful bijection.

The tolerance here is exactly zero, and that is worth a remark. Most checks in this series compare floating-point quantities and must state a tolerance, as in fd01 §10. This one compares the same stored values reached by two different routes, so no arithmetic is performed and any discrepancy at all would be a bookkeeping error, not a rounding error. Knowing which kind of check you are making tells you which tolerance is honest.

Solving the problem: greedy versus exhaustive¶

Now the economics. A matching assigns each worker to a distinct firm; we hold it as a set of pairs, and its total surplus is $\sum_{(x,y)} \Phi_{xy}$ over that set.

The obvious heuristic is greedy: repeatedly take the largest available surplus, remove that worker and that firm, and continue. It is fast, it is intuitive, and it is what a manager filling posts one at a time would do.

In [36]:
def total_surplus(matching):
    # matching: a set of (x, y) pairs
    return sum(Phi_x_y[xy] for xy in matching)

def greedy_matching():
    # repeatedly take the largest remaining surplus among unmatched workers and firms
    unmatched_x, unmatched_y = set(X), set(Y)
    matching = set()
    while unmatched_x:
        best = max(((x, y) for x in unmatched_x for y in unmatched_y),
                   key=lambda xy: Phi_x_y[xy])
        matching.add(best)
        unmatched_x.discard(best[0])
        unmatched_y.discard(best[1])
    return matching

greedy = greedy_matching()
print("greedy matching:")
for x, y in sorted(greedy):
    print(f"  {x:<10} -> {y:<12} surplus {Phi_x_y[(x, y)]:>5.1f}")
print(f"  total surplus: {total_surplus(greedy):.1f}")
greedy matching:
  analyst    -> bank         surplus  10.0
  engineer   -> startup      surplus   6.0
  manager    -> university   surplus   8.0
  total surplus: 24.0

With $n_x = 3$ and $n_y = 4$ we can afford to check every possibility: there are $4 \cdot 3 \cdot 2 = 24$ ways to assign three distinct firms to three workers. itertools.permutations enumerates them.

In [37]:
from itertools import permutations

candidates = [set(zip(X, assignment)) for assignment in permutations(Y, nbx)]
best = max(candidates, key=total_surplus)

print(f"enumerated {len(candidates)} feasible matchings\n")
print("optimal matching:")
for x, y in sorted(best):
    print(f"  {x:<10} -> {y:<12} surplus {Phi_x_y[(x, y)]:>5.1f}")
print(f"  total surplus: {total_surplus(best):.1f}")

print(f"\ngreedy  = {total_surplus(greedy):.1f}")
print(f"optimal = {total_surplus(best):.1f}")
print(f"greedy loses {total_surplus(best) - total_surplus(greedy):.1f} "
      f"({1 - total_surplus(greedy) / total_surplus(best):.1%} of the maximum)")
enumerated 24 feasible matchings

optimal matching:
  analyst    -> startup      surplus   9.0
  engineer   -> bank         surplus   9.0
  manager    -> university   surplus   8.0
  total surplus: 26.0

greedy  = 24.0
optimal = 26.0
greedy loses 2.0 (7.7% of the maximum)

Reading the result¶

Greedy fails, and the way it fails is the point.

Greedy opens by giving the bank to the analyst, because $\Phi_{\text{analyst},\text{bank}} = 10$ is the largest single entry in the table. The analyst does indeed have an absolute advantage at the bank. But the engineer is nearly as good there ($9$) and poor everywhere else, whereas the analyst is nearly as good at the startup ($9$) as at the bank. Handing the bank to the analyst therefore strands the engineer in a job worth $6$, and the extra surplus captured on the first move is more than given back on the second.

The optimal assignment sends the engineer to the bank and the analyst to the startup. What governs it is not who is best at a job but who is best relative to their next-best alternative: comparative advantage, in the sense that has organized the theory of trade and assignment since Ricardo. A rule that looks only at the level of $\Phi_{xy}$, one match at a time, cannot see this, because comparative advantage is a statement about differences of $\Phi$ across the whole table.

Two threads run forward from here. First, the reason the correct answer is hard to compute is that the choices interact: the number of feasible matchings is $n_y!/(n_y - n_x)!$, which is $24$ here but astronomical at realistic sizes, so enumeration is not an algorithm. lp01 and ot01 show that the problem is nonetheless a linear program, solvable in polynomial time. Second, and more interesting economically, that linear program has a dual, and its dual variables are wages: a vector $u_x$ of worker values and $v_y$ of firm values with $u_x + v_y \geq \Phi_{xy}$ for every pair, holding with equality exactly on the matched pairs. Comparative advantage is precisely what a price system decentralizes: the market implements the optimal assignment through wages, without anyone enumerating anything. That is the subject of ot01, and the discount factors of fd01 §11 were the first, one-dimensional instance of the same idea.

13. PEP 8 and the Pythonic idiom¶

PEP 8 is the style guide. There is no need to memorize it, but the highlights:

  • snake_case for variables, functions, and modules; CapWords for classes; ALL_CAPS for module-level constants.
  • Four spaces per indent level; lines under 79 characters, or 88 under the Black and ruff conventions we adopt in fd04.
  • Imports at the top, one per line, grouped: standard library, then third-party, then local.
  • A docstring on every public function (fd03).
  • Two blank lines between top-level definitions, one inside a class.

A handful of taste-level idioms make code feel Pythonic:

Less Pythonic More Pythonic
if len(xs) > 0: if xs:
if x == None: if x is None:
for i in range(len(xs)): use(xs[i]) for x in xs: use(x)
for i in range(len(xs)): use(i, xs[i]) for i, x in enumerate(xs): use(i, x)
result = []
for x in xs: result.append(f(x))
result = [f(x) for x in xs]
d.has_key(k) (Python 2) k in d

To this list the 'math+econ+code' series adds the naming convention introduced in fd01 §6 and used for the first time in §12 above: an array's name records its axes, Phi_x_y for a surplus indexed by $(x,y)$, Phi_xy for its row-major flattening, and nbx, nby for the dimensions. In fd04 we install ruff and black to enforce PEP 8 automatically. The axis convention no tool can check for you; it is worth the discipline anyway, because it is the one that prevents the errors a linter cannot see.

14. Summary¶

  • Four containers, four jobs. Lists for sequences you will modify, tuples for fixed records and, crucially, for composite keys, dicts for lookup in $O(1)$, sets for membership and set algebra. Choosing badly is not merely inelegant: a membership test against a list inside a loop is the most common avoidable performance bug in research code.

  • Assignment binds names to objects; it never copies. Mutating an object is visible through every name bound to it, and a mutable default argument is evaluated once and shared across calls. The failure mode is not a crash but a plausible wrong number, which is why this section is the important one.

  • Comprehensions, enumerate, zip, and starred unpacking are not decoration. They remove the index arithmetic that off-by-one errors live in, and zip(..., strict=True) turns a silently truncated result into an exception.

  • The lecture's real object is the index map of §12: a dictionary from a pair $(x,y)$ to a position in a flat vector, verified as a bijection by three independent routes. It is the explicit form of the row-major flattening $\operatorname{vec}_C$ that fd08 performs with reshape(-1) and that the ot and lp series use to assemble constraint matrices. Writing it out once by hand is what makes the one-line version auditable later.

  • The economics is in §12's last table. The greedy rule assigns the bank to the worker with the highest surplus there and loses nearly 8% of the attainable total, because the optimal assignment is governed by comparative rather than absolute advantage: a property of differences across the whole surplus table, invisible to any rule that looks at one match at a time. In ot01 the same problem is solved as a linear program whose dual variables are wages: the price system decentralizes comparative advantage, and no one has to enumerate anything.

15. Exercises¶

Write your answer in the cell below each prompt. Worked solutions are in §17, at the end of this notebook; attempt each exercise before reading them.

Exercise 1: Growth rates for a panel. Given the nominal-GDP panel

gdp_panel = {"USA":     [21.4, 21.5, 23.6, 25.46, 27.7],
             "Germany": [3.89, 3.94, 4.27, 4.07, 4.46],
             "Japan":   [5.12, 5.06, 5.00, 4.23, 4.21]}

(2019–2023, USD trillions, rounded), build a dict mapping each country to its list of year-on-year growth rates, using a dict comprehension with a nested list comprehension.

Then check your work against a second route: for each country, verify that compounding the growth rates from the 2019 level reproduces the 2023 level, reporting the largest relative gap across the three countries and the tolerance you accept. Which country's series is not monotone, and what happened to it?

In [38]:
# your answer here

Exercise 2: Inverting a dict. Given gdp_2022, build the reverse mapping from value to country.

What goes wrong if two countries have identical GDP? Demonstrate the failure by constructing a small dict where it occurs, then write a version returning a dict[float, list[str]] that survives collisions. State, in one sentence, the condition on the original dict under which inversion to a single-valued mapping is well defined, and note that this is exactly the condition for a function to be invertible.

In [39]:
# your answer here

Exercise 3: A panel as a dict of tuples, and its index map. Using the data of Exercise 1, build a dict gdp_c_t keyed by (country, year). Then, following §12:

  1. build the index map idx_c_t sending each (country, year) to its row-major position, with year varying fastest;
  2. produce the flat vector gdp_ct in that order;
  3. verify the bijection as in §12: the ordering from the map must equal the ordering from the nested loops, and the inverse map must recover every key;
  4. write a one-line comprehension returning all entries for "USA", and a second returning all entries for 2022.

Which of the two one-liners scans the whole dictionary, and how would you avoid that if the panel had a million entries?

In [40]:
# your answer here

Exercise 4: A prime sieve as a generator, and a small proof. Write a generator primes_below(n) yielding the primes strictly below $n$ by trial division, and test it with list(primes_below(50)).

Efficiency proper comes in fd04, but one optimization is free and worth proving: it suffices to test divisors $d$ with $d \leq \sqrt{m}$. Prove it, show that if $m$ is composite then it has a divisor $d$ with $1 < d \le \sqrt m$, and implement it. Verify your generator against an independent route by checking that it produces exactly the known 25 primes below 100, and report the count of trial divisions saved by the $\sqrt{m}$ bound at $n = 2000$.

In [41]:
# your answer here

Exercise 5: Spot the bug. The function below is meant to return a list of $n$ empty bins, each of which is itself an empty list. It contains a reference-aliasing bug of the kind discussed in §7.

def make_bins(n):
    return [[]] * n

bins = make_bins(3)
bins[0].append("first")
print(bins)

Predict what the buggy version prints, then run it and explain the result in terms of names and objects. Give the one-line fix, and verify it with an id()-based test that would fail on the buggy version and pass on the fixed one.

Finally: a matching in §12 was held as a set of pairs. If instead you had built a list of per-worker bins with make_bins, this bug would have silently merged every worker's assignments into one. What would the reported total surplus have looked like: too high, too low, or plausible?

In [42]:
# your answer here

16. Further directions¶

You can now express most everyday data manipulation in idiomatic Python without reaching for a library, and you have built the index map that the rest of the series will keep reusing.

fd03 is about organizing code: functions with default and keyword arguments, where the mutable-default trap of §7 becomes a design question rather than a curiosity, *args and **kwargs, scope, closures, modules and packages, and a first look at classes and @dataclass. After that, fd04 measures the claim of §6, Exercise 4 there times list membership against set membership and finds a factor of $10^5$ at a million elements, and puts this kind of code under test. fd05 brings Git and the research compendium. Then the scientific stack proper: in fd08 the dict Phi_x_y of §12 becomes a NumPy array, the flattening becomes reshape(-1), and the index map becomes the thing that lets a Kronecker product assemble a constraint matrix.

Save your work, restart the kernel, and run all cells top-to-bottom before you move on.

17. Solutions to the exercises¶

Reference solutions, using only what this lecture and fd01 introduce. As in the lecture, each numerical answer is checked against a second route with a stated tolerance.

Solution to Exercise 1: Growth rates for a panel¶

In [43]:
gdp_panel = {"USA":     [21.4, 21.5, 23.6, 25.46, 27.7],
             "Germany": [3.89, 3.94, 4.27, 4.07, 4.46],
             "Japan":   [5.12, 5.06, 5.00, 4.23, 4.21]}

growth_panel = {country: [(levels[t] - levels[t - 1]) / levels[t - 1]
                          for t in range(1, len(levels))]
                for country, levels in gdp_panel.items()}

for country, rates in growth_panel.items():
    print(f"{country:<9}", "  ".join(f"{g:+7.2%}" for g in rates))
USA        +0.47%   +9.77%   +7.88%   +8.80%
Germany    +1.29%   +8.38%   -4.68%   +9.58%
Japan      -1.17%   -1.19%  -15.40%   -0.47%
In [44]:
# check: compounding the growth rates from the 2019 level must give the 2023 level
print(f"{'country':<9}{'compounded':>14}{'reported':>12}{'rel. gap':>12}")
print("-" * 47)
worst = 0.0
for country, levels in gdp_panel.items():
    compounded = levels[0]
    for g in growth_panel[country]:
        compounded *= (1 + g)
    rel = abs(compounded - levels[-1]) / levels[-1]
    worst = max(worst, rel)
    print(f"{country:<9}{compounded:>14.6f}{levels[-1]:>12.2f}{rel:>12.1e}")

tol = 1e-12
print("-" * 47)
print(f"worst relative gap = {worst:.1e}   (tolerance {tol:.0e})")
assert worst < tol, "the growth rates do not compound back to the reported levels"
print("check passed.")
country      compounded    reported    rel. gap
-----------------------------------------------
USA           27.700000       27.70     0.0e+00
Germany        4.460000        4.46     0.0e+00
Japan          4.210000        4.21     0.0e+00
-----------------------------------------------
worst relative gap = 0.0e+00   (tolerance 1e-12)
check passed.

The check is a genuine one despite looking circular: the growth rates were built by differencing, and compounding them back is the inverse operation, so agreement confirms that the index arithmetic in the comprehension, the range(1, len(levels)) and the t-1, is right. That is precisely where an off-by-one would hide.

Japan is the non-monotone series: it falls in every year but the first, from $5.12T in 2019 to $4.21T in 2023, an 18% decline. Nothing collapsed in Japan; the series is in nominal US dollars, and the yen depreciated sharply against the dollar over 2022–2023. Japanese GDP measured in yen rose over the same period. This is not a subtlety about Python: it is the most common way that a correctly computed growth rate answers a different question from the one intended, and no amount of verification inside the code would catch it. Verification checks that you computed what you said; only economics checks that you said the right thing.

Solution to Exercise 2: Inverting a dict¶

In [45]:
inverse_gdp = {value: country for country, value in gdp_2022.items()}
print("inverted:", inverse_gdp)
print("lookup by value:", inverse_gdp[25.46])
print(f"sizes: original {len(gdp_2022)}, inverted {len(inverse_gdp)}")
inverted: {17.96: 'China', 2.78: 'France', 4.07: 'Germany', 4.23: 'Japan', 3.07: 'UK', 25.46: 'USA'}
lookup by value: USA
sizes: original 6, inverted 6
In [46]:
# the failure: two countries with identical GDP -- one silently overwrites the other
tied = {"Japan": 4.2, "Germany": 4.2, "UK": 3.1}
naive = {value: country for country, value in tied.items()}

print("original:", tied)
print("inverted:", naive, " <- Japan has vanished")
print(f"sizes: {len(tied)} keys in, {len(naive)} keys out\n")

# the robust version: map each value to the LIST of countries attaining it
robust = {}
for country, value in tied.items():
    robust.setdefault(value, []).append(country)

print("robust inversion:", robust)
assert sum(len(v) for v in robust.values()) == len(tied), "entries were lost"
print("check passed: every original entry is preserved.")
original: {'Japan': 4.2, 'Germany': 4.2, 'UK': 3.1}
inverted: {4.2: 'Germany', 3.1: 'UK'}  <- Japan has vanished
sizes: 3 keys in, 2 keys out

robust inversion: {4.2: ['Japan', 'Germany'], 3.1: ['UK']}
check passed: every original entry is preserved.

setdefault(value, []) returns the list already stored under value, inserting a fresh empty list first if there is none: the standard idiom for accumulating into a dict of lists. (collections.defaultdict(list) does the same thing with less ceremony.)

The condition. Inversion to a single-valued mapping is well defined exactly when the original dict is injective, no two keys share a value, which is the same condition under which a function admits an inverse. When it fails, the dict comprehension does not complain: later keys silently overwrite earlier ones, and the only visible symptom is that the result has fewer entries than the input. Comparing len before and after is the cheap check, and it is the one worth writing down, because a mapping that silently loses rows is exactly the sort of bug that survives to publication.

Solution to Exercise 3: A panel as a dict of tuples, and its index map¶

In [47]:
years = [2019, 2020, 2021, 2022, 2023]
countries_panel = list(gdp_panel)
nbc, nbyr = len(countries_panel), len(years)

# the panel, keyed by the composite index (country, year)
gdp_c_t = {(c, year): gdp_panel[c][t]
           for c in countries_panel for t, year in enumerate(years)}

# the index map: row-major, with the year varying fastest
idx_c_t = {(c, year): i * nbyr + t
           for i, c in enumerate(countries_panel) for t, year in enumerate(years)}

# the flat vector in that order
gdp_ct = [gdp_c_t[key] for key in sorted(idx_c_t, key=idx_c_t.get)]

print(f"nbc = {nbc}, nbyr = {nbyr}, flat length = {len(gdp_ct)} = {nbc * nbyr}")
print("first eight entries:", gdp_ct[:8])
nbc = 3, nbyr = 5, flat length = 15 = 15
first eight entries: [21.4, 21.5, 23.6, 25.46, 27.7, 3.89, 3.94, 4.27]
In [48]:
# verification, exactly as in section 12
order_from_map = sorted(idx_c_t, key=idx_c_t.get)
order_from_loops = [(c, year) for c in countries_panel for year in years]
assert order_from_map == order_from_loops, "the index map is not row-major"

inverse = {position: key for key, position in idx_c_t.items()}
assert all(inverse[idx_c_t[k]] == k for k in idx_c_t), "the round trip fails"

gap = max(abs(gdp_ct[idx_c_t[k]] - gdp_c_t[k]) for k in gdp_c_t)
print(f"max |flat[idx(c,t)] - panel[c,t]| = {gap:.1e}   (exact; tolerance 0)")
assert gap == 0.0
print("check passed: the panel flattening is a faithful bijection.")
max |flat[idx(c,t)] - panel[c,t]| = 0.0e+00   (exact; tolerance 0)
check passed: the panel flattening is a faithful bijection.
In [49]:
usa_entries = {k: v for k, v in gdp_c_t.items() if k[0] == "USA"}
year_2022 = {k: v for k, v in gdp_c_t.items() if k[1] == 2022}

print("USA  :", usa_entries)
print("2022 :", year_2022)
USA  : {('USA', 2019): 21.4, ('USA', 2020): 21.5, ('USA', 2021): 23.6, ('USA', 2022): 25.46, ('USA', 2023): 27.7}
2022 : {('USA', 2022): 25.46, ('Germany', 2022): 4.07, ('Japan', 2022): 4.23}

Which one scans everything? Both do. Each comprehension walks all $n_c \times n_{yr}$ items and filters, so each costs $O(n_c n_{yr})$: fine for fifteen entries, wasteful for a million.

The fix is to stop asking a flat dict a question it is not indexed for. Three options, in increasing order of seriousness. For the country query, the keys are known in advance, so build the lookup directly rather than scanning: {(c, y): gdp_c_t[(c, y)] for y in years} costs $O(n_{yr})$. More generally, maintain a secondary index, a dict from country to the list of its keys, which is what a database calls an index and what fd06's pandas MultiIndex builds for you. And when both axes are dense, as here, the right answer is not a dict at all but the two-dimensional array of fd08, where selecting a row or a column is a slice rather than a search.

That progression, dict of tuples, then secondary index, then array, is the same one the series follows over the next five lectures, and the reason is visible here: the dict is the right structure while the object is sparse and irregular, and the wrong one as soon as it is dense and rectangular.

Solution to Exercise 4: A prime sieve as a generator, and a small proof¶

Claim. If $m > 1$ is composite, it has a divisor $d$ with $1 < d \le \sqrt m$.

Proof. Since $m$ is composite it factors as $m = ab$ with $1 < a \le b < m$. If $a > \sqrt m$ then $b \ge a > \sqrt m$, so $ab > \sqrt m \cdot \sqrt m = m$, contradicting $ab = m$. Hence $a \le \sqrt m$, and $d = a$ is the required divisor. $\blacksquare$

Contrapositive, which is the form the algorithm uses: if no $d$ with $1 < d \le \sqrt m$ divides $m$, then $m$ is prime. So trial division may stop at $\sqrt m$ rather than at $m - 1$.

In [50]:
def is_prime(m):
    if m < 2:
        return False
    d = 2
    while d * d <= m:          # equivalent to d <= sqrt(m), without a square root
        if m % d == 0:
            return False
        d += 1
    return True

def primes_below(n):
    for m in range(2, n):
        if is_prime(m):
            yield m

print(list(primes_below(50)))
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]
In [51]:
# independent check: the count of primes below 100 is known to be 25
primes_100 = list(primes_below(100))
print(f"primes below 100: {len(primes_100)} found, 25 expected")
assert len(primes_100) == 25, "the sieve does not reproduce the known count"

# and every one of them must be prime by the definition, tested directly
assert all(all(p % d != 0 for d in range(2, p)) for p in primes_100)
print("check passed: all 25 are prime by exhaustive trial division.")
primes below 100: 25 found, 25 expected
check passed: all 25 are prime by exhaustive trial division.
In [52]:
# how much work does the sqrt bound save?
def trial_divisions(n, bounded):
    count = 0
    for m in range(2, n):
        d = 2
        while (d * d <= m) if bounded else (d < m):
            count += 1
            if m % d == 0:
                break
            d += 1
    return count

n = 2000
with_bound = trial_divisions(n, bounded=True)
without = trial_divisions(n, bounded=False)

print(f"trial divisions below n = {n}")
print(f"  d < m       : {without:>9,}")
print(f"  d <= sqrt(m): {with_bound:>9,}")
print(f"  saved       : {without - with_bound:>9,}  ({1 - with_bound / without:.1%})")
trial divisions below n = 2000
  d < m       :   281,802
  d <= sqrt(m):    13,397
  saved       :   268,405  (95.2%)

The bound removes about 95% of the divisions at $n = 2000$, and the saving grows with $n$: testing one number $m$ falls from $O(m)$ to $O(\sqrt m)$. It costs one line and follows from three lines of proof, which is the general shape of the best optimizations, and the opposite of the "make it faster by rewriting it in a faster language" reflex we examine in fd04.

Note also what the generator buys here. primes_below never builds a list, so it can be consumed lazily, next(primes_below(10**9)) is instantaneous, and it could equally well be written unbounded, yielding primes forever. A function that returns a list cannot do that.

Solution to Exercise 5: Spot the bug¶

In [53]:
def make_bins_buggy(n):
    return [[]] * n

bins = make_bins_buggy(3)
bins[0].append("first")
print("buggy :", bins)
buggy : [['first'], ['first'], ['first']]

It prints [['first'], ['first'], ['first']].

The explanation is §7 exactly. The expression [[]] builds a list containing one empty list object. The * 3 then builds a new outer list containing three references: to that same inner object, not to three copies of it. There is only ever one inner list, reached by three routes, so appending through any one of them is visible through all three. Nothing was copied, because in Python nothing is ever copied unless you ask.

The fix is to evaluate [] afresh for each bin, which a comprehension does by construction: return [[] for _ in range(n)].

In [54]:
def make_bins(n):
    return [[] for _ in range(n)]          # a fresh list per bin

# an id()-based test: the bins must be distinct objects
for name, factory in (("buggy", make_bins_buggy), ("fixed", make_bins)):
    bins = factory(3)
    distinct = len({id(b) for b in bins})
    print(f"{name:<6}: {distinct} distinct inner objects out of 3")

bins = make_bins(3)
assert len({id(b) for b in bins}) == 3, "the bins are aliased"
bins[0].append("first")
print("\nfixed :", bins)
print("check passed: appending to one bin leaves the others empty.")
buggy : 1 distinct inner objects out of 3
fixed : 3 distinct inner objects out of 3

fixed : [['first'], [], []]
check passed: appending to one bin leaves the others empty.

The consequence for §12. Had the matching been held in per-worker bins built this way, every worker's assignment would have landed in the same list. The bins would all report the full set of matches, so summing $\Phi$ over "each worker's" bin would count every matched pair once per worker: inflating the total by a factor of $n_x$ and reporting a surplus of about $78$ where the truth is $26$.

The number would be too high, and, this is the part worth internalizing, it would still be deterministic, reproducible, and quietly proportional to the right answer. Re-running the notebook would reproduce it exactly. A colleague checking your code by re-running it would confirm it. Only a check against an independent route, or against the economics, a total surplus exceeding the sum of the largest entry in each row is impossible, would expose it. That is why fd01 insisted on a second route for every number, and why the aliasing section of this lecture is the important one.