Nested Loops

Just as the body of a function can contain an if clause or a loop, so too can the body of a loop contain another entire loop.

for i in range(2):
    # The outer loop body, including the entire next loop, is repeated twice.
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0
1 1
1 2

A loop within another loop is called a nested loop. Nested loops are useful for iterating over combinations of values from multiple collections.

Double Loops

In a nested loop structure with one outer loop and one inner loop, the entire inner loop is completed once for each outer loop iteration. Step through the preceding example, focusing first just on the outer loop.

1

for i in range(2):
    for j in range(3):
        print(i, j)

The outer loop variable i is assigned to the first value of range(2): i = 0.

2

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2

The body of the outer loop is executed with i = 0.

3

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2

The outer loop advances to the next value of range(2): i = 1.

4

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0
1 1
1 2

The body of the outer loop is executed with i = 1.

The following steps through this example one line of execution at a time.

1

for i in range(2):
    for j in range(3):
        print(i, j)

The outer loop variable i is assigned to the first value of range(2): i = 0.

2

for i in range(2):
    for j in range(3):
        print(i, j)

The inner loop variable j is assigned to the first value of range(3): j = 0. The outer loop variable is still i = 0.

3

for i in range(2):
    for j in range(3):
        print(i, j)

0 0

The loop body runs with i = 0 and j = 0.

4

for i in range(2):
    for j in range(3):
        print(i, j)

0 0

The inner loop advances: j is reassigned to the second value of range(3), j = 1. The outer loop variable is still i = 0.

5

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1

The loop body runs with i = 0 and j = 1.

6

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1

The inner loop advances again: j is reassigned to the third value of range(3), j = 2. The outer loop variable is still i = 0.

7

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2

The loop body runs with i = 0 and j = 2.

8

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2

The inner loop has now produced every value of range(3), so it ends. Execution returns to the outer loop, which advances: i is reassigned to the second value of range(2), i = 1.

9

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2

The inner loop starts over from scratch: j is reassigned to 0.

10

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0

The loop body runs with i = 1 and j = 0.

11

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0

The inner loop advances: j is reassigned to 1. The outer loop variable is still i = 1.

12

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0
1 1

The loop body runs with i = 1 and j = 1.

13

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0
1 1

The inner loop advances again: j is reassigned to 2. The outer loop variable is still i = 1.

14

for i in range(2):
    for j in range(3):
        print(i, j)

0 0
0 1
0 2
1 0
1 1
1 2

The loop body runs with i = 1 and j = 2.

The inner loop is exhausted again, and now the outer loop has also produced every value of range(2), so it ends too. Execution continues after both loops.

In this example the outer loop runs twice, and for each of those the inner loop runs three times, so the inner loop body executes \(3 \times 2 = 6\) times in total. In general, the total number of passes through the innermost line is the product of the lengths of the loops: two loops that each run \(n\) times produce \(n \times n = n^2\) passes. Because of this, nested loops that each have a large number of iterations can slow down a program significantly.

Counting the total number of innermost iterations is not always straightforward because the bounds of an inner loop can change each time the outer loop progresses.

Quick Check

What is the output of the following code?

total = 0
for i in range(8):
    for j in range(i):
        total += 1
print(total)

Check Answer

28

• The outer loop is over range(8), assigning i = 0, 1, 2, ..., 7.

• The inner loop is over range(i), which has i iterations. The first time through the outer loop there are zero iterations, then one, then two, and so on to seven.

• The total is therefore \(0 + 1 + 2 + \dots + 7 = 28\).

Problem 3

A lattice point is a point \((x, y)\) whose coordinates are both integers. The circle of radius \(r = 2.5\) shown below encloses 21 lattice points.

Write a function named lattice_points() that accepts a nonnegative number \(r\) and counts the lattice points lying on or within the circle of radius \(r\) centered at the origin. That is, find the number of pairs \((x, y)\) satisfying \( x^2 + y^2 \le r^2 \) where \(x\) and \(y\) are both integers. Return the number of lattice points.

Hint: The radius \(r\) may not be an integer, but range() only accepts integers. A built-in function can help with this.

Test Cases

print(lattice_points(0))
print(lattice_points(1))
print(lattice_points(2.5))
print(lattice_points(20))

1
5
21
1257

Another natural use of a double loop is combining every entry of one collection with every entry of another.

Problem 4

As we saw in Problem 1, a polynomial \(p(x) = a_0 + a_1 x + a_2 x^2 + a_3 x^3 + \cdots\) can be represented with a list of coefficients in ascending order of degree, so that element i is the coefficient \(a_i\) multiplying \(x^{i}\).

• \(p(x) = 1 + 5x + 3x^2 + 9x^3\) can be represented with [1, 5, 3, 9].

• \(p(x) = x^2\) can be represented with [0, 0, 1].

• \(p(x) = 4\) can be represented with [4].

The product of a polynomial \(p(x)\) with \(n\) coefficients and a polynomial \(q(x)\) with \(m\) coefficients is a polynomial with \(n + m - 1\) coefficients:

\[\begin{split} \begin{aligned} p(x) &= a_0 + a_1 x + a_2 x^2 + \cdots + a_{n-1} x^{n-1} = \sum_{i=0}^{n-1} a_i x^i, \\ q(x) &= b_0 + b_1 x + b_2 x^2 + \cdots + b_{m-1} x^{m-1} = \sum_{j=0}^{m-1} b_j x^j, \\ p(x)q(x) &= a_0 b_0 + (a_0 b_1 + a_1 b_0)x + (a_0 b_2 + a_1 b_1 + a_2 b_0)x^2 + \cdots = \sum_{i=0}^{n-1}\sum_{j=0}^{m-1} a_i b_j x^{i + j}. \end{aligned} \end{split}\]

Another way to describe this is that the coefficient on \(x^{k}\) in the product polynomial \(p(x)q(x)\) is the sum of all products \(a_i b_j\) such that \(i + j = k\).

Write a function named poly_product() that accepts a NumPy array of \(n\) coefficients representing a polynomial \(p(x)\) and a NumPy array of \(m\) coefficients representing a polynomial \(q(x)\). Return a NumPy array of the \(n + m - 1\) coefficients of the product polynomial \(p(x) q(x)\).

Hint: enumerate() may be helpful.

Test Cases

# (1 + x)(1 - x) = 1 - x^2
print(poly_product(np.array([1, 1]), np.array([1, -1])))

# (2 + 3x + x^2)(1 - 4x) = 2 - 5x - 11x^2 - 4x^3
print(poly_product(np.array([2, 3, 1]), np.array([1, -4])))

# (1 + x + x^2 + x^3)(0) = 0 + 0x + 0x^2 + 0x^3
print(poly_product(np.array([1, 1, 1, 1]), np.array([0])))

[ 1.  0. -1.]
[  2.  -5. -11.  -4.]
[0. 0. 0. 0.]

Triple Loops

Nesting can extend past two levels: a loop inside a loop inside a loop examines every combination of three collections.

count = 0
for a in range(2):
    for b in range(3):
        for c in range(4):
            count += 1
print(count)

24

Here, the innermost line runs \(2 \times 3 \times 4 = 24\) times. In general, \(k\) nested loops that each run \(n\) times produce \(n^{k}\) passes through the innermost code. The exponent grows with the number of nested loops, and small increases in the exponent lead to enormous increases in work. For \(n = 1000\), \(n^2\) is only a million iterations, which is fast for a computer. But \(n^3\) is a billion iterations, noticeably slows a program, and \(n^7\) is a sextillion iterations, which is not realistic to run.

Sometimes a triple or quadruple loop appears to be necessary, but it can be avoided by rephrasing part of the problem. Removing one level of nesting drastically reduces the number of iterations, which can turn a slow program into one that finishes quickly.

Problem 5

A Pythagorean triplet is a set of three positive integers \((a, b, c)\) satisfying \( a^2 + b^2 = c^2. \) For example, \(3^2 + 4^2 = 5^2\), so \((3, 4, 5)\) is a Pythagorean triplet.

Write a function named pythagorean_triplet() that accepts a positive integer \(n\) and finds a Pythagorean triplet \((a, b, c)\) such that \( a + b + c = n. \) Return the three values \(a\), \(b\), and \(c\); if there are no such values, the function should return None. Only return values forming a single triplet, even if there are multiple possible triplets for a given \(n\).

Your function should be able to complete the test cases within a few seconds.

Hint: Once \(a\) and \(b\) are chosen, do you need another loop for \(c\)?

Test Cases

print(pythagorean_triplet(12))
print(pythagorean_triplet(13))
print(pythagorean_triplet(150))
print(pythagorean_triplet(10050))

(3, 4, 5)
None
(25, 60, 65)
(1072, 4425, 4553)