Solving problems with loops step by step

Many small programming problems are loops in disguise: is this number prime, what is the greatest common divisor of two numbers, what do the digits of a number add up to. Each needs only a few lines of Python, and the hard part is rarely the syntax. It is deciding what the loop should do on each pass and when it should stop. This lesson plans that first, in plain words, then writes the code, then measures how the work grows as the numbers get bigger.

Plan the loop before you write it

Take a small problem: add up the decimal digits of a whole number, so 9475 gives 9 + 4 + 7 + 5 = 25. Before writing any code, answer four questions:

1. What goes in and what comes out? A whole number n that is 0 or more; the sum of its digits.
2. What changes on each pass? One digit moves out of n and into a running total: n % 10 is the last digit and n //= 10 drops it.
3. What stays true at the start of every pass? total plus the digit sum of what is left in n always equals the answer. A statement like this, true before and after every pass, is called the loop’s invariant.
4. When does it stop? When n is 0, no digits are left, so by the invariant total is the answer.

The code follows from the plan line by line. The program also traces the loop for 9475 and checks the invariant at the top of each pass:

def digit_sum(n):
    """Add up the decimal digits of a whole number n that is 0 or more."""
    total = 0
    while n > 0:
        total += n % 10  # add the last digit
        n //= 10  # then drop it
    return total


# Trace the loop for 9475, and check the invariant at the top of every pass.
n, total = 9475, 0
while n > 0:
    print(f"n = {n:<5} total = {total:<3} total + digit_sum(n) = {total + digit_sum(n)}")
    total += n % 10
    n //= 10
print(f"n = {n:<5} total = {total:<3} n is 0, so the loop stops")

print(digit_sum(9475), digit_sum(7), digit_sum(0), digit_sum(-25))

n = 9475  total = 0   total + digit_sum(n) = 25
n = 947   total = 5   total + digit_sum(n) = 25
n = 94    total = 12  total + digit_sum(n) = 25
n = 9     total = 16  total + digit_sum(n) = 25
n = 0     total = 25  n is 0, so the loop stops
25 7 0 0

The last column stays 25 on every line: that is the invariant at work, and seeing it hold is a quick way to trust a loop. Two edge cases are worth checking for every loop: the smallest inputs and the ones the plan did not expect. For 0 the loop never runs and the answer, 0, is right. For -25 it also never runs, and 0 is wrong. The plan said “0 or more”, so either refuse negative numbers with an error or work with abs(n).

Primes by trial division

A prime is a whole number above 1 whose only divisors are 1 and itself. The direct test, trial division, tries the possible divisors one by one and stops at the first one that divides exactly, because one divisor is enough to prove the number is not prime. The lesson on while loops wrote this search as while divisor * divisor <= n; here it is with for and range(), together with the reason why it never needs to try anything above the square root of n:

Divisors come in pairs whose product is n, and in each pair the smaller one is at most the square root. Any number that is not prime therefore has a divisor small enough for the loop to reach. math.isqrt(n) gives the integer square root, the largest whole number whose square is at most n:

import math


def is_prime(n):
    """True when the whole number n is a prime."""
    if n < 2:
        return False
    for d in range(2, math.isqrt(n) + 1):
        if n % d == 0:
            return False  # a divisor: n is not prime, stop at once
    return True


def is_prime_wrong(n):
    """The same with a common mistake: range() stops one short of isqrt(n)."""
    if n < 2:
        return False
    for d in range(2, math.isqrt(n)):
        if n % d == 0:
            return False
    return True


print("primes below 30:", [n for n in range(30) if is_prime(n)])
for n in [0, 1, 2, 4, 9, 15, 25, 97]:
    print(f"{n:>2}: is_prime {is_prime(n)!s:<5}  is_prime_wrong {is_prime_wrong(n)}")

# Going through a float can be off: one too large for the first number, one too small for the second.
for n in [1_000_000_007**2 - 1, 9_509_733_067_886_049**2]:
    print(f"int(math.sqrt(n)) = {int(math.sqrt(n))}, math.isqrt(n) = {math.isqrt(n)}")

primes below 30: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
 0: is_prime False  is_prime_wrong False
 1: is_prime False  is_prime_wrong False
 2: is_prime True   is_prime_wrong True
 4: is_prime False  is_prime_wrong True
 9: is_prime False  is_prime_wrong True
15: is_prime False  is_prime_wrong True
25: is_prime False  is_prime_wrong True
97: is_prime True   is_prime_wrong True
int(math.sqrt(n)) = 1000000007, math.isqrt(n) = 1000000006
int(math.sqrt(n)) = 9509733067886048, math.isqrt(n) = 9509733067886049

• n < 2 is checked first: 0 and 1 are not prime, and the loop would wrongly accept them.
• For 2 and 3, range(2, math.isqrt(n) + 1) is empty, so the loop does not run and the function returns True.
• The + 1 matters. is_prime_wrong stops one divisor short, a common off-by-one mistake, so it never tries isqrt(n) itself. It calls 4 a prime, because for 4 to 8 its loop is empty; it accepts 9 and 25, the squares of primes, whose only small divisor is the square root; and it accepts 15, whose smallest divisor, 3, is isqrt(15).
• math.isqrt() works with whole numbers and is always exact. int(math.sqrt(n)) goes through a floating-point number, which keeps only 53 bits of precision, and the last two lines show it coming out one too large for one big number and one too small for another. Too small is the dangerous one: for the square of a prime, the loop would stop just before the only divisor that proves it is not prime.

Euclid’s algorithm for the greatest common divisor

The greatest common divisor (gcd) of two whole numbers is the largest number that divides both. Euclid’s algorithm finds it with one simple step, repeated: replace the pair (a, b) with (b, a % b), until b is 0. The plan is short. The invariant is that the pair always has the same gcd as the original numbers, because a number divides both a and b exactly when it divides both b and a % b; the loop stops when b is 0, and then the gcd is a. This version prints every step for the pair 1071 and 462:

import math


def gcd(a, b):
    """Euclid's algorithm: replace (a, b) with (b, a % b) until b is 0."""
    while b != 0:
        print(f"{a} = {a // b} × {b} + {a % b}")
        a, b = b, a % b
    return a


print("gcd(1071, 462) =", gcd(1071, 462))
print("math.gcd(1071, 462) =", math.gcd(1071, 462))
print("math.lcm(1071, 462) =", math.lcm(1071, 462))

1071 = 2 × 462 + 147
462 = 3 × 147 + 21
147 = 7 × 21 + 0
gcd(1071, 462) = 21
math.gcd(1071, 462) = 21
math.lcm(1071, 462) = 23562

Three passes are enough here, because the remainders shrink fast. In real code, use the standard library: math.gcd() takes any number of arguments (since Python 3.9; before that, exactly two) and math.lcm(), new in Python 3.9, returns the least common multiple. For two numbers it is their product divided by their gcd: 1071 × 462 ÷ 21 = 23,562.

Digits again: Armstrong numbers

An Armstrong number (also called a narcissistic number) equals the sum of its digits, each raised to the power of the number of digits. The digit loop from the start of the lesson needs one change, the power:

def is_armstrong(n):
    """True when n equals the sum of its digits, each raised to the power of the number of digits."""
    power = len(str(n))
    total = 0
    rest = n
    while rest > 0:
        total += (rest % 10) ** power
        rest //= 10
    return total == n


print([n for n in range(100, 1000) if is_armstrong(n)])
print("153 =", " + ".join(f"{d}³" for d in "153"), "=", 1**3 + 5**3 + 3**3)

[153, 370, 371, 407]
153 = 1³ + 5³ + 3³ = 153

There are exactly four with three digits. The plan is the same as for the digit sum; only what is added on each pass has changed, which is why it pays to learn the pattern rather than each program.

How running time grows

Before you time anything, you can count. For a prime n, trial division tries every divisor it is allowed, so the count depends only on where the loop stops. This program counts the divisions for the largest prime below 100, 1,000, 10,000, 100,000 and 1,000,000:

import math


def is_prime(n):
    if n < 2:
        return False
    for d in range(2, math.isqrt(n) + 1):
        if n % d == 0:
            return False
    return True


def divisions(n, last):
    """How many divisions trial division makes for n, trying the divisors 2, 3, ..., last."""
    count = 0
    for d in range(2, last + 1):
        count += 1
        if n % d == 0:
            break
    return count


print(f"{'prime n':>9} {'up to n - 1':>12} {'up to isqrt(n)':>15}")
for limit in [100, 1_000, 10_000, 100_000, 1_000_000]:
    n = limit - 1
    while not is_prime(n):  # the largest prime below the limit
        n -= 1
    print(f"{n:>9,} {divisions(n, n - 1):>12,} {divisions(n, math.isqrt(n)):>15,}")

  prime n  up to n - 1  up to isqrt(n)
       97           95               8
      997          995              30
    9,973        9,971              98
   99,991       99,989             315
  999,983      999,981             998

Each time the number grows ten times, the loop up to n - 1 makes ten times as many divisions, but the loop up to isqrt(n) only about three times as many, since the square root of 10 is about 3.16. The usual shorthand for this is big-O notation: the first loop is O(n) and the second O(√n), meaning the work grows at most in proportion to n and to the square root of n.

The table counts each division as one step, which is fair for numbers of everyday size. Python’s integers have unlimited precision, though, and dividing a number with thousands of digits takes longer than dividing a small one, so for huge numbers the time grows faster than the count of steps.

Counts predict, and a clock confirms. This program times both versions on 999,983, the largest prime below a million:

import math
import time


def is_prime_slow(n):
    """Trial division with every divisor up to n - 1."""
    if n < 2:
        return False
    for d in range(2, n):
        if n % d == 0:
            return False
    return True


def is_prime(n):
    """Trial division with divisors up to isqrt(n) only."""
    if n < 2:
        return False
    for d in range(2, math.isqrt(n) + 1):
        if n % d == 0:
            return False
    return True


for check in [is_prime_slow, is_prime]:
    start = time.perf_counter()
    answer = check(999_983)
    milliseconds = (time.perf_counter() - start) * 1000
    print(f"{check.__name__}(999_983) = {answer} in {milliseconds:.3f} ms")

is_prime_slow(999_983) = True in 33.358 ms
is_prime(999_983) = True in 0.042 ms

The recorded times come from the computer named under the output, and yours will differ, in the browser more so, because Python runs more slowly there. What does not change is the shape: a thousand times fewer divisions makes the second version hundreds of times faster. Counting the work first tells you whether a loop will be fast enough before you wait for it.

Clear first, fast later

Write the clearest loop you can while you are solving a problem, test it on small inputs whose answers you know, and only then make it faster or shorter. Keep the tests, so that each change can be checked against them. For real programs, prefer the standard library: math.isqrt(), math.gcd() and math.lcm() work exactly with whole numbers of any size. And when you need many primes rather than one, a different algorithm, the sieve of Eratosthenes in this lesson’s exercise, does far less work than calling is_prime() again and again.

Key takeaways

• Plan a loop in words first: what goes in and comes out, what changes on each pass, what stays true (the invariant), and when it stops.
• Test the smallest inputs and the ones the plan excludes, such as 0, 1 and negative numbers.
• Trial division needs divisors only up to math.isqrt(n), and stops at the first one that divides; range()’s stop value needs the + 1.
• Euclid’s algorithm repeats a, b = b, a % b until b is 0; real code calls math.gcd() and math.lcm().
• Count operations to see how work grows: up to n is O(n), up to √n is O(√n), and timing confirms what the count predicts.