Loop invariants with Python assert

Why this matters

A loop can give the right answer on one test and the wrong answer on another. A beginner often runs the program once, sees a plausible number, and stops. That check tells you about one input. It does not tell you why the loop works.

A loop invariant is a short statement about the variables that is true every time the loop is about to start a pass. If you can write that statement, and you can show it stays true from one pass to the next, you have a proof that the loop does what you want. The proof works for every input, not only the one you tried.

This lesson uses two small jobs, adding up a list and finding the largest value in a list. Both are simple enough to prove completely. The earlier lessons in this module showed how a for loop runs. Here the question is what the variables must mean at each step.

Words used in this lesson

A pass is one trip through the body of the loop.

An invariant is a statement that does not change its truth from one pass to the next. The word means "does not vary". In a loop, the statement is about the variables and is checked at one fixed point, which is the start of each pass, before the body runs.

A prefix is the first part of a list. The prefix of length i is the items at positions 0 up to i - 1. Writing values[0..i-1] means that prefix. In Python, values[:i] is the same prefix. Position 0 is the first item, because Python counts from zero.

The first example: a running sum

Here is a loop that adds the marks in a list and prints the variables before and after each pass. The variable is called total instead of sum, because sum is the name of a built-in function. Using that name for a variable hides the function until the program ends.

The program prints:

before i=0 sum=0
after  i=0 sum=4
before i=1 sum=4
after  i=1 sum=13
before i=2 sum=13
after  i=2 sum=15
before i=3 sum=15
after  i=3 sum=22
final sum=22

The header range(len(marks)) produces the positions 0, 1, 2 and 3, one at a time. It stops by itself, so there is no test to write in the header.

Look at the lines that say before. Before the pass with i=1, the total is 4, which is the total of marks[0] alone. Before the pass with i=2, the total is 13, which is 4 + 9. The pattern is clear: before pass i, the total is the total of the items before position i.

That pattern is the invariant. Written in plain words:

Before each pass with index i, total equals the total of marks[0..i-1].

When i is 0, the range 0..-1 is empty and the total of nothing is 0. That matches the first before line. In Python, sum(marks[:i]) computes the same total directly.

Proving the sum with three steps

A proof for a loop has three parts. The same three parts work for almost every loop you will write.

1. The start. Before the first pass, the invariant must hold. Here i is 0 and total is 0, and the total of an empty range is 0. The invariant holds.

2. The step. If the invariant holds before a pass, it must hold before the next pass. Suppose it holds before a pass with index i. The body adds marks[i] to total, so total now equals the total of marks[0..i-1] plus marks[i]. That is the total of marks[0..i]. The next pass uses position i + 1, and the total of marks[0..(i+1)-1] is exactly that. So the invariant holds before the next pass.

3. The end. The loop stops when range has no more positions. The last pass was for position len(marks) - 1. The invariant then says total equals the total of every item in the list. So the final value of total is correct.

You do not need to run the code to get this result. The proof uses only the body, the start and the point where the loop stops.

Checking the invariant while the program runs

Python has a statement called assert. It checks a condition. If the condition is false, Python stops the program with an AssertionError and prints the message after the comma. Python turns asserts on by default. Java is the other way round, where asserts stay off until you turn them on with -ea. To turn asserts off in Python, run the program with python3 -O, where -O means "optimize", and Python skips every assert statement.

The next program checks the invariant before every pass. A helper function adds up the prefix directly, so the check does not depend on the loop it is testing.

Run it with python3 check_sum.py. The output is:

sum = 14

The helper is slow, because it adds the prefix again on every pass. That is fine for a test. In normal code, you keep the check while you are learning or testing, and you remove it or run with python3 -O in the finished program.

Seeing a broken invariant

Now change one line. Instead of adding values[i] to total, the bad version sets total to values[i]. The line is total = values[i] in the file below. Everything else is the same.

def sum_of_prefix(values, end):
    total = 0
    for k in range(end):
        total += values[k]
    return total
 
values = [5, -2, 8, 3]
total = 0
for i in range(len(values)):
    assert total == sum_of_prefix(values, i), f"invariant broken at i={i}"
    total = values[i]
print("sum =", total)

Run it with python3 check_sum_bad.py. Python prints a traceback, which is a report of where the error happened. The file name and line number at the top of the traceback depend on where you save the file, so they are left out here. The last lines are:

    assert total == sum_of_prefix(values, i), f"invariant broken at i={i}"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: invariant broken at i=2

The check at the start of pass i=2 failed. Before that pass, the bad version has total equal to -2, but the prefix total is 5 + -2, which is 3. The first two passes happened to agree, which is why the bug did not show up at the start. The invariant check catches the error at the first pass where the two values differ.

Run the same file with python3 -O check_sum_bad.py and the assert does nothing. The program prints sum = 3 and reports no error. This is why the check needs to be on while you test.

The second example: a running maximum

The second job is to find the largest value in a list. The invariant is about the largest value seen so far.

Before each pass with index i, best equals the largest value in scores[0..i-1].

Here is the code:

The output is:

good max = 9

Python has a built-in function max that does this job. The loop is here to show the proof.

The proof has the same three parts.

The start. The loop begins with i equal to 1 and best equal to scores[0]. The prefix scores[0..0] has one item, and its largest value is scores[0]. The invariant holds.

The step. Suppose best is the largest of scores[0..i-1] before a pass. If scores[i] is larger than best, the body sets best to scores[i]. That value is larger than everything before it, so best is now the largest of scores[0..i]. If scores[i] is not larger, best is kept. The item scores[i] is no larger than best, and best was already the largest of the earlier items, so best is still the largest of scores[0..i]. Either way, the invariant holds before the next pass.

The end. The loop stops when range has no more positions. The invariant then says best is the largest of the whole list.

Where the invariant tells you where the start must be

The start is an easy place to go wrong. Some people write best = 0 because it looks like a neutral starting value. Here is that version next to the correct one, run on the same kinds of input.

def max_bad(scores):
    best = 0
    for i in range(len(scores)):
        if scores[i] > best:
            best = scores[i]
    return best
 
def max_good(scores):
    best = scores[0]
    for i in range(1, len(scores)):
        if scores[i] > best:
            best = scores[i]
    return best
 
scores = [4, 9, 2, 7, 9, 3]
print("good max =", max_good(scores))
negatives = [-5, -2, -9]
print("bad max  =", max_bad(negatives))
print("good max =", max_good(negatives))
empty = []
print(max_good(empty))

The output is:

good max = 9
bad max  = 0
good max = -2
Traceback (most recent call last):
  ...
    best = scores[0]
           ~~~~~~^^^
IndexError: list index out of range

The bad max line reports 0 for a list in which every value is negative. The largest value in [-5, -2, -9] is -2, so 0 is wrong. The invariant explains why. At the start, best is supposed to be the largest of the empty prefix when i is 0, but the empty prefix has no largest value, so the invariant cannot be made true at the start. Setting best to 0 only pretends it is true. Starting from scores[0] is correct because that item is a real member of the prefix.

The last line shows the empty list. max_good reads scores[0] before it checks the length, so an empty list fails. The message list index out of range says the same thing: position 0 does not exist. If empty input is possible, the function should check if not scores: first and decide what to return.

The third example: the position of the largest value

Finding a position is harder than finding a value, because ties matter. Suppose two items are both the largest. The loop must decide which position to report, and the invariant must say which one it keeps.

The invariant here is about best_index, the position of the largest value so far. Before each pass with index i, values[best_index] is the largest value in values[0..i-1].

The output is:

index 1 holds 8

The value 8 appears at positions 1 and 3. The loop reports 1, because the test is >, which is strict. A value at position 3 equal to 8 does not replace position 1. If you change the test to >=, the loop reports position 3 instead. Both answers are correct, and the invariant must say which rule you chose. That is why the comment in the code states the tie rule. A proof that says only "the largest value" would not be enough to decide between the two answers.

The expression values.index(max(values)) also gives 1 here. The index method returns the first position where a value appears.

A harder job: the average

An average needs two running values, the total and the count. The invariant must describe both, and both must agree at every pass.

The average of a list is the total divided by the number of items. Here the invariant says that before a pass with index i, total is the sum of values[0..i-1], and count is i. The second part is simple, because count goes up by 1 on every pass. The first part is the same as the running sum from earlier.

The output is:

good average = 5.5

At the end, the loop has visited every position, so total is the sum of every item and count is the number of items. In Python, / always gives a decimal number, even when the answer is whole. That is why the division here gives 5.5 with no extra step.

The mistake to watch for in Python is the // operator. It is floor division. It keeps the whole part of the division and drops the fraction. The next program shows it. It is the same loop, but the last line uses //.

The output is:

bad average  = 1
good average = 1.5

The total is 3 and the count is 2. 3 // 2 is 1, so the function returns the whole number 1, which prints as 1. The loop itself was correct. The mistake was in the choice of division operator. The invariant says what the totals are, but it does not say which division the final line uses, so this kind of error is not caught by the proof about the loop. Check the last line separately.

Common mistakes

A common mistake is to write an invariant with the wrong upper bound. If the invariant says the total equals the total of values[0..i], it is false before the first pass, when i is 0. Here is the assert from the first example with i + 1 written in place of i:

def sum_of_prefix(values, end):
    total = 0
    for k in range(end):
        total += values[k]
    return total
 
values = [5, -2, 8, 3]
total = 0
for i in range(len(values)):
    assert total == sum_of_prefix(values, i + 1), f"invariant broken at i={i}"
    total += values[i]
print("sum =", total)

The output is:

    assert total == sum_of_prefix(values, i + 1), f"invariant broken at i={i}"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: invariant broken at i=0

The check fails on the first pass. That is a useful result, because it tells you the statement is wrong, not the loop. Before the pass, total is 0 and the total of values[0..0] is 5. Fix the statement so that it matches the state before the body runs.

A second mistake is to check only the final output. A loop can produce the right final answer for one input and still be wrong for others. The invariant checks each pass, so a bug shows up in the pass where it starts, not at the end.

A third mistake is to forget that asserts can be turned off. In the broken-sum example, python3 -O prints sum = 3 and nothing else. The program looks finished, but the check that would have caught the bug never ran. Run without -O whenever you test a loop that uses asserts.

A fourth mistake is to put parentheses around an assert's condition and message. Python then reads the two parts as one tuple, and a tuple that is not empty is always true. The assert never fails.

The output is:

SyntaxWarning: assertion is always true, perhaps remove parentheses?
  assert (total == 1, "total should be 1")
tuple assert passed

The first line also starts with the file name and line number, which are left out here. The warning names the problem. Write assert total == 1, "total should be 1" without the parentheses.

A fifth mistake is to put the invariant in a comment and never check it. A comment is a claim, and claims can be wrong. Either prove the claim on paper, as in the three steps above, or check it with an assert while you test.

Practice

Exercise 1. Write a loop that counts the even numbers in the list values. State the invariant before a pass with index i, and say what it means at the end.

Show the answer

The code is:

count = 0
for i in range(len(values)):
    if values[i] % 2 == 0:
        count += 1

The invariant is: before a pass with index i, count is the number of even values in values[0..i-1]. It starts at 0, which is correct for an empty prefix. Each pass adds 1 only when values[i] is even, so the invariant holds before the next pass. At the end, the loop has visited every position, so count is the number of even values in the whole list.

Exercise 2. Write a loop that finds the smallest value in a non-empty list. State its invariant, and say why the start must use the first value rather than 0.

Show the answer

The code is:

smallest = values[0]
for i in range(1, len(values)):
    if values[i] < smallest:
        smallest = values[i]

The invariant is: before a pass with index i, smallest is the least value in values[0..i-1]. The start must use values[0] because 0 is not guaranteed to be in the list. If every value is positive, starting from 0 would report 0, which is not in the list at all. Python also has min(values), which gives the same answer.

Exercise 3. Trace the loop from exercise 2 on [6, -1, 4]. Write the value of i and smallest before each pass.

Show the answer

Before the pass with i=1, smallest is 6. Before the pass with i=2, smallest is -1, because -1 is less than 6. The pass with i=2 compares 4 with -1, and 4 is not smaller, so smallest stays -1. At the end, smallest is -1, which is the least value in the list.

Exercise 4. The loop below is meant to add the numbers from 1 to n. Write the invariant that the loop needs, and show that it holds at the start.

total = 0
for i in range(1, n + 1):
    total += i

Show the answer

The invariant is: before a pass with value i, total is the sum of 1 through i-1. At the start, the first value is 1, and the sum of no numbers is 0, so total is 0 and the invariant holds. Adding i to total gives the sum of 1 through i, and the next pass has value i+1, so the invariant holds again. When the loop stops, every value from 1 through n has been added, so total is the sum of 1 through n.

Exercise 5. Write a loop that computes n factorial, which is 1 times 2 times 3 up to n. Give the invariant, and check it for n equal to 5 and for n equal to 0.

Show the answer

The code is:

product = 1
for i in range(1, n + 1):
    product *= i

The invariant is: before a pass with value i, product is the product of 1 through i-1. For n equal to 5, the result is 120. For n equal to 0, the range is empty, the body never runs, product stays 1, and 1 is the correct value of 0 factorial. The standard library function math.factorial(5) also gives 120.

Summary

• An invariant is a statement about the variables that is true before every pass of a loop.
• A proof has three parts: the statement is true at the start, it stays true after each pass, and it gives the right answer when the loop stops.
• For a running sum, the invariant says the total equals the total of the items before the current index.
• For a running maximum, the start must use a real item such as scores[0], not a guessed value such as 0.
• Use assert to check an invariant while you test, and run with python3 -O only when you want the checks off.
• Ties need a rule in the invariant, because the loop must say which of two equal values it keeps.