Problem Set 3: NumPy

Homework for Module 3

Introduction

Use these exercises to practise the concepts from this module. I recommend deactivating Gemini first, along with any other AI helper, and attempting the exercises on your own.

Mandatory exercises are highlighted in red boxes. These are the exercises you must complete to receive your participation points. You may use AI help, but remember that you may be asked to explain your code in class.

The remaining exercises are optional. Try each one before opening its collapsed solution.

When an exercise asks you to predict an output or error, explain your reasoning before checking the answer. For programming exercises, test your solution with additional inputs.

import numpy as np

Introduction to NumPy

Shadowing the NumPy Alias

Predict what happens when this code runs. Why does np.array() fail?

np = 4
arr = np.array([2, 3, np, 5])

The assignment np = 4 replaces the NumPy alias with an integer. An integer has no array attribute, so Python raises an AttributeError. Restore the alias before using NumPy again:

import numpy as np
arr = np.array([2, 3, 4, 5])
print(arr)
[2 3 4 5]
Replacing a Sum Loop

Which NumPy expression replaces the loop and returns one total?

arr = np.array([1, 2, 3, 4, 5])

total = 0
for value in arr:
    total = total + value

Choose among np.cumsum(arr), np.sum(arr), and np.loop(arr).

total = np.sum(arr)
print(total)
15

np.cumsum(arr) returns every running total, while np.sum(arr) returns one value. NumPy has no np.loop() function.

Array Dimensions

Which expression reports the number of dimensions of this array?

arr = np.array([[1, 2, 3], [4, 5, 6]])

Choose among arr.shape, arr.ndim, arr.ndim(), and np.arr.ndim.

arr.ndim is the integer 2. It is a property, so it does not use parentheses. arr.shape is (2, 3) and answers a different question.

Image Shape

A color image uses a three-dimensional array. Which assignment retrieves its height and width without storing the color-channel count?

img = np.array([
    [[1, 2, 3], [4, 5, 6]],
    [[7, 8, 9], [10, 11, 12]]
])
height, width, _ = img.shape
print(height, width)
2 2

The shape is (height, width, channels). The underscore records that the third value is intentionally unused.

Empty Arrays

Which expressions correctly create an empty NumPy array?

# Option A
# x = np.array()

# Option B
x = np.array([])

# Option C
y = np.zeros(0, dtype=int)

Options B and C work and produce arrays with shape (0,). Option A fails because np.array() requires an object argument.

Array-Creation Debugging

Determine which lines work. For each valid line, give the output shape and data type.

a = np.zeros(3).astype(float)
c = np.ones(5, dtype=bool)
d = np.ones((5, 5))

# Deliberately incorrect
# b = np.zeros(3, 3)
  • a has shape (3,) and a floating-point dtype.
  • c has shape (5,) and Boolean dtype.
  • d has shape (5, 5) and a floating-point dtype.
  • np.zeros(3, 3) is invalid because the shape must be supplied as one tuple for a 2D array: np.zeros((3, 3)).
Battleship

Represent a \(10 \times 10\) Battleship board with a NumPy array, where 0 is water and 1 is part of a ship. Using only np.zeros(), indexing, and slicing (no loops):

  1. Place a horizontal ship of length 4 in row 2, starting at column 3.
  2. Place a vertical ship of length 3 in column 8, starting at row 5.
  3. Place a one-cell ship in the bottom-left corner, using negative indices.
  4. Your opponent shoots at row 2, column 5. Print the value of the board at that cell: is it a hit?
  5. Print the shape and size of the board, and the whole row 2.
board = np.zeros((10, 10), dtype=int)

board[2, 3:7] = 1     # horizontal ship: columns 3, 4, 5, 6
board[5:8, 8] = 1     # vertical ship: rows 5, 6, 7
board[-1, 0] = 1      # bottom-left corner

print(board[2, 5])    # 1, a hit!
print(board.shape, board.size)
print(board[2])
print(board)
1
(10, 10) 100
[0 0 0 1 1 1 1 0 0 0]
[[0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 0]
 [0 0 0 1 1 1 1 0 0 0]
 [0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 1 0]
 [0 0 0 0 0 0 0 0 1 0]
 [0 0 0 0 0 0 0 0 1 0]
 [0 0 0 0 0 0 0 0 0 0]
 [1 0 0 0 0 0 0 0 0 0]]

Using dtype=int keeps the board readable when printed. Slices place the whole ship in one assignment, which with nested lists would require a loop.

Exam Versions

To avoid copying, no two students sitting next to each other (left, right, in front, or behind) may get the same version of the exam. Using only np.zeros(), np.arange(), reshape(), indexing, and slicing (no loops):

  1. Write a function named seating_plan(n_rows, n_cols) that returns a 2D array with the seat numbers of the room, starting at 1 and numbered row by row.
  2. Write a function named exam_versions(n_rows, n_cols) that returns a 2D integer array with the version of each seat: 0 for version A and 1 for version B. The seat in the first row and first column gets version A.

Sample input 1: seating_plan(3, 4)
Expected result 1:

[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]

Sample input 2: exam_versions(3, 4)
Expected result 2:

[[0 1 0 1]
 [1 0 1 0]
 [0 1 0 1]]
def seating_plan(n_rows, n_cols):
    return np.arange(1, n_rows * n_cols + 1).reshape(n_rows, n_cols)


def exam_versions(n_rows, n_cols):
    versions = np.zeros((n_rows, n_cols), dtype=int)
    versions[::2, 1::2] = 1   # even rows, odd columns
    versions[1::2, ::2] = 1   # odd rows, even columns
    return versions


print(seating_plan(3, 4))
print(exam_versions(3, 4))
[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]
[[0 1 0 1]
 [1 0 1 0]
 [0 1 0 1]]
Holiday Lights

Inspired by Advent of Code 2015, Day 6.

You have a grid of lights, all of them off (0). Each instruction switches on (1) or off (0) a whole rectangle of lights, given by two opposite corners (r1, c1) and (r2, c2), both included.

Write a Python function switch(grid, action, r1, c1, r2, c2) that applies one instruction to the grid, where action is either "on" or "off". Do not use loops.

Sample input:

grid = np.zeros((6, 6), dtype=int)
switch(grid, "on", 0, 0, 3, 3)
switch(grid, "off", 1, 1, 2, 2)
switch(grid, "on", 2, 4, 5, 5)
print(grid)

Expected result:

[[1 1 1 1 0 0]
 [1 0 0 1 0 0]
 [1 0 0 1 1 1]
 [1 1 1 1 1 1]
 [0 0 0 0 1 1]
 [0 0 0 0 1 1]]

Array Arithmetic and Broadcasting

Matrix Multiplication

Which expressions perform matrix multiplication for these two-dimensional arrays?

a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6], [7, 8]])

Choose among a * b, a.b, a @ b, and np.matmul(a, b).

Both a @ b and np.matmul(a, b) perform matrix multiplication. The expression a * b multiplies element by element, while a.b is not valid syntax for this operation.

Assigning One Value to a Row

Predict the final value of x.

x = np.array([[1, 2], [3, 4]])
x[0] = 0
print(x)
[[0 0]
 [3 4]]

The scalar 0 broadcasts across the first row:

[[0 0]
 [3 4]]
Fixing Vector Addition

The following code should add [1, 2, 3] and [4, 5, 6]. Fix both array-creation errors.

x = np.array(1, 2, 3)
y = np.arange(4, 6, 1)
print(x + y)
x = np.array([1, 2, 3])
y = np.arange(4, 7)
print(x + y)
[5 7 9]

np.array() receives one collection, and the stopping value of np.arange() is excluded.

Broadcasting a Row Vector

Predict the shape and values of c.

a = np.array([1, 2, 3])
b = np.array([[4, 5, 6], [7, 8, 9]])
c = b > a

The shape is (2, 3). NumPy compares a with every row of b:

[[ True  True  True]
 [ True  True  True]]
Incompatible Shapes

Explain why this comparison fails.

a = np.array([1, 2, 3])
b = np.array([4, 5, 6, 7])
c = b > a

The trailing dimensions are 4 and 3. They are neither equal nor 1, so the arrays are not broadcast-compatible.

Row and Column Assignment

Predict the two outputs.

rows = np.array([[1, 2], [3, 4]])
rows[1, :] = 0

columns = np.array([[1, 2], [3, 4]])
columns[:, 0] = 0

print(rows)
print(columns)
[[1 2]
 [0 0]]
[[0 2]
 [0 4]]
[[1 2]
 [0 0]]

[[0 2]
 [0 4]]

The colon selects every position along the corresponding axis.

Final Grades with Any Weights

Write a function named final_grades(marks, weights) that receives a 2D array of marks between 0 and 1 (one row per student, one column per exercise) and a 1D array with the points of each exercise. It must return the final grade of each student out of 10, without loops.

The points do not necessarily add up to 100, so divide by the total number of points of the exam.

Sample input:

marks = np.array([[1, 0.5, 0], [0.5, 1, 1]])
weights = np.array([50, 30, 20])
final_grades(marks, weights)

Expected result: [6.5 7.5]

def final_grades(marks, weights):
    return 10 * (marks @ weights) / np.sum(weights)


marks = np.array([[1, 0.5, 0], [0.5, 1, 1]])
weights = np.array([50, 30, 20])
print(final_grades(marks, weights))
[6.5 7.5]

marks @ weights gives the points of each student, and dividing by the total points of the exam turns them into a fraction of the maximum.

Factorial, Revisited

In Module 1 you computed factorials with a loop. Now do it without one:

  1. Write a function factorial(n) that uses np.arange() and np.prod().
  2. Create an array with the factorials of 1 to 10 in a single line, using np.cumprod().
  3. Compute factorial(21). Is the result correct? Why?

Sample input: factorial(5)
Expected result: 120

def factorial(n):
    return np.prod(np.arange(1, n + 1))


print(factorial(5))
print(np.cumprod(np.arange(1, 11)))
print(factorial(21))
120
[      1       2       6      24     120     720    5040   40320  362880
 3628800]
-4249290049419214848

\(21! \approx 5.1 \cdot 10^{19}\) does not fit in a 64-bit integer (the maximum is about \(9.2 \cdot 10^{18}\)), so the result overflows and NumPy returns a wrong (even negative) number without any error. Python integers have no such limit, which is why math.factorial(21) is correct. With floats, np.prod(np.arange(1, n + 1, dtype=float)) gives an approximation instead.

Polynomial Operation, Revisited

In Module 1 you evaluated these functions one value of x at a time, with \(a = 1\), \(b = 1\), \(c = -6\):

\[f(x) = a \cdot x^2+b \cdot x+c \qquad g(x) = \frac{a}{x} \qquad h(x) = \frac{f(x)}{g(x)}\]

  1. Store the points \(x = -2, 0, 2.1\) in one array and compute \(f\), \(g\), and \(h\) for all of them at once, without loops.
  2. Plain Python raised a ZeroDivisionError at \(x = 0\). What does NumPy do instead? What are the values of \(g(0)\) and \(h(0)\)?
  3. Compute the simplified version, \(h(x) = x^3 + \frac{b}{a} x^2 + \frac{c}{a} x\), on the same array. Do both versions agree at every point?
  4. Evaluate the simplified \(h\) on 101 evenly spaced points between \(-3\) and \(3\), and find the smallest value and the x where it happens.

Expected result for steps 1 and 2: f = [-4. -6. 0.51], g = [-0.5 inf 0.476...], h = [8. -0. 1.071], with a warning about a division by zero.

a, b, c = 1, 1, -6
x = np.array([-2, 0, 2.1])

f = a * x**2 + b * x + c
g = a / x
h = f / g
print(f, g, h)

h_simplified = x**3 + (b / a) * x**2 + (c / a) * x
print(h_simplified)

x_grid = np.linspace(-3, 3, 101)
h_grid = x_grid**3 + (b / a) * x_grid**2 + (c / a) * x_grid
i = np.argmin(h_grid)
print(h_grid[i], x_grid[i])
[-4.   -6.    0.51] [-0.5                inf  0.47619048] [ 8.    -0.     1.071]
[8.    0.    1.071]
-4.0588560000000005 1.1399999999999997
/tmp/ipykernel_3215/3553276906.py:5: RuntimeWarning: divide by zero encountered in divide
  g = a / x

NumPy does not raise an error when dividing by zero: it shows a RuntimeWarning and returns inf. Then \(h(0) = -6 / \infty = -0\), which matches the simplified version, \(h(0) = 0\). Both versions agree at every point.

Squares, Sums, and Powers

Inspired by Project Euler, Problem 6 and Problem 29.

The sum of the squares of the first ten natural numbers is \(1^2 + 2^2 + \dots + 10^2 = 385\), while the square of their sum is \((1 + 2 + \dots + 10)^2 = 3025\). The difference between them is \(3025 - 385 = 2640\).

Write a Python function sum_square_difference(n) that computes this difference for the first n natural numbers, without loops.

Now consider all the powers \(a^b\) with \(2 \le a \le 5\) and \(2 \le b \le 5\). Some of them are repeated, such as \(2^4 = 4^2 = 16\), so there are only 15 distinct values. Write a Python function distinct_powers(n) that counts the distinct values of \(a^b\) for \(2 \le a \le n\) and \(2 \le b \le n\), without loops. Why does your function fail for n = 100?

Sample input 1: sum_square_difference(10)
Expected result 1: 2640

Sample input 2: distinct_powers(5)
Expected result 2: 15


Loading and Saving Data

Text and Binary Formats

Write code that saves the array below once as readable text and once in NumPy’s binary format. Then load both files and compare their shapes.

measurements = np.array([[1.2, 1.4], [1.8, 2.1]])
np.savetxt("measurements.txt", measurements)
np.save("measurements.npy", measurements)

from_text = np.loadtxt("measurements.txt")
from_binary = np.load("measurements.npy")

print(from_text.shape)
print(from_binary.shape)

Both loaded arrays have shape (2, 2). The .npy format also preserves NumPy-specific dtype and shape information exactly.

Join Daily Recordings

A sensor saves one text file per day, with one value per line. Write a function named join_recordings(filenames) that receives a list of file names, loads each file with np.loadtxt(), and returns a single 1D array with all the values in the order of the list.

Then save the joined array as a binary .npy file, load it back, and check that it is identical to the original.

Sample input:

np.savetxt("day1.txt", [1, 2, 3])
np.savetxt("day2.txt", [4, 5])
join_recordings(["day1.txt", "day2.txt"])

Expected result: [1. 2. 3. 4. 5.]

def join_recordings(filenames):
    arrays = []
    for filename in filenames:
        arrays.append(np.loadtxt(filename))
    return np.concatenate(arrays)


np.savetxt("day1.txt", [1, 2, 3])
np.savetxt("day2.txt", [4, 5])
joined = join_recordings(["day1.txt", "day2.txt"])
print(joined)

np.save("joined.npy", joined)
print(np.array_equal(np.load("joined.npy"), joined))  # True

Array Manipulation

Joining Two Vectors

Join a and b into one one-dimensional array.

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
c = np.concatenate([a, b])
print(c)
[1 2 3 4 5 6]

Adding the arrays would produce [5, 7, 9]; it would not join them.

Predicting a Reshape

What are the shape and values of reshaped?

arr = np.array([1, 2, 3, 4, 5, 6])
reshaped = arr.reshape(2, 3)

The result has two rows and three columns:

[[1 2 3]
 [4 5 6]]
Vertical Stacking

Stack a = [1, 2, 3] and b = [4, 5, 6] as two rows of one matrix.

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
stacked = np.vstack([a, b])
print(stacked)
[[1 2 3]
 [4 5 6]]

The result has shape (2, 3).

Sliding Windows

Fix the code so it creates

[[ 1,  2,  3,  4],
 [ 2,  3,  4,  5],
 ...
 [11, 12, 13, 14]]

from z = np.arange(1, 15).

z = np.arange(1, 15)
windows = []

for index in range(1, len(z) - 3):
    windows.append(z[index:index + 3])

result = np.array(windows)
z = np.arange(1, 15)
windows = []

for index in range(len(z) - 3):
    windows.append(z[index:index + 4])

result = np.array(windows)
print(result)
[[ 1  2  3  4]
 [ 2  3  4  5]
 [ 3  4  5  6]
 [ 4  5  6  7]
 [ 5  6  7  8]
 [ 6  7  8  9]
 [ 7  8  9 10]
 [ 8  9 10 11]
 [ 9 10 11 12]
 [10 11 12 13]
 [11 12 13 14]]

The first window starts at index 0, and each slice must contain four values.

reshape() Does Not Modify in Place

Fix the code so it creates the requested matrix.

x = np.linspace(1, 9, 9)
x.reshape((3, 3))
x[1, 1] = 0
x = np.linspace(1, 9, 9)
x = x.reshape((3, 3))
x[1, 1] = 0
print(x)
[[1. 2. 3.]
 [4. 0. 6.]
 [7. 8. 9.]]

reshape() returns the reshaped array. Assign that return value to a variable before using two indices.

Fibonacci Array Debugging

The function should return the first n Fibonacci values as an integer array. Fix the initialization and return statement.

def fibonacci_sequence(n: int):
    arr = np.array()

    for index in range(n):
        if index <= 1:
            new_element = 1
        else:
            new_element = arr[-2] + arr[-1]
        arr = np.append(arr, new_element)

    return int(arr)
def fibonacci_sequence(n: int) -> np.ndarray:
    arr = np.array([], dtype=int)

    for index in range(n):
        if index <= 1:
            new_element = 1
        else:
            new_element = arr[-2] + arr[-1]
        arr = np.append(arr, new_element)

    return arr


print(fibonacci_sequence(10))
[ 1  1  2  3  5  8 13 21 34 55]

np.array([]) creates an empty array; int(arr) incorrectly tries to turn the complete array into one integer.

Concatenating an Empty Array

Predict the output.

x = np.array([])
y = np.zeros(3)
a = np.concatenate([x, y])
print(a)
[0. 0. 0.]

Python prints [0. 0. 0.]. Concatenating an empty array contributes no elements, and np.zeros(3) uses a floating-point dtype by default.

Points per Part

An exam is divided into parts, each made of consecutive exercises. Write a function named part_points(marks, weights, cuts) that receives:

  • marks: a 2D array of marks between 0 and 1, one row per student and one column per exercise.
  • weights: a 1D array with the points of each exercise.
  • cuts: a list with the column indices where each new part starts.

It must return a 2D array with one row per student and one column per part, containing the points of each student in each part. Use the splitting and stacking functions of this session.

Sample input:

marks = np.array([[1, 0.5, 1, 0, 1], [0.5, 0.5, 0, 1, 1]])
weights = np.array([20, 20, 10, 10, 40])
part_points(marks, weights, [2, 4])

Expected result:

[[30. 10. 40.]
 [20. 10. 40.]]
def part_points(marks, weights, cuts):
    mark_parts = np.hsplit(marks, cuts)
    weight_parts = np.split(weights, cuts)

    columns = []
    for m, w in zip(mark_parts, weight_parts):
        columns.append(np.sum(m * w, axis=1))
    return np.column_stack(columns)


marks = np.array([[1, 0.5, 1, 0, 1], [0.5, 0.5, 0, 1, 1]])
weights = np.array([20, 20, 10, 10, 40])
print(part_points(marks, weights, [2, 4]))
[[30. 10. 40.]
 [20. 10. 40.]]
Largest Product in a Series

Inspired by Project Euler, Problem 8.

In the series of digits "123456789", the three adjacent digits with the largest product are 7, 8, and 9, with \(7 \cdot 8 \cdot 9 = 504\).

Write a Python function largest_product(digits, k) that receives a string of digits and returns the largest product of k adjacent digits, together with those digits. Do not loop over the positions of the series. Hint: you can build a 2D array where each row contains k adjacent digits.

Sample input 1: largest_product("123456789", 3)
Expected result 1: (504, "789")

Sample input 2: largest_product("3141592653589793238462643383279", 4)
Expected result 2: (4536, "8979")


Masking

Mask Length

Explain why this code fails, then repair it.

arr = np.arange(4)
arr_bool = np.array([True, False, True, False, True])
selected = arr[arr_bool]

A Boolean mask must have the same length as the axis it selects. Either extend arr to five elements or shorten the mask:

arr = np.arange(4)
arr_bool = np.array([True, False, True, False])
selected = arr[arr_bool]
print(selected)
[0 2]
Filtering Values

Without a loop, select every value less than or equal to 3.

arr = np.array([1, 2, 3, 4, 5])
result = arr[arr <= 3]
print(result)
[1 2 3]
Selecting Even Values

Predict the output.

arr = np.array([1, 2, 3, 4, 5])
mask = arr % 2 == 0
print(arr[mask])
[2 4]

Python prints [2 4]. The mask is True exactly where the remainder after division by 2 is zero.

Replacing Selected Values

Replace every value greater than 3 with 10.

arr = np.array([1, 2, 3, 4, 5])
arr[arr > 3] = 10
print(arr)
[ 1  2  3 10 10]

Assigning to a loop variable would only rebind that variable; it would not modify the array.

Selecting Odd Values

Write one Boolean expression that selects all odd values from the array.

arr = np.array([1, 2, 3, 4, 5, 6])
odds = arr[arr % 2 != 0]
print(odds)
[1 3 5]
Combining Two Conditions

The code should return the positive odd values [1, 3], but its second mask has the wrong length. Fix it.

arr = np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4])
mask_odd = arr % 2 != 0
mask_positive = arr[mask_odd] > 0
result = arr[mask_positive]

Build both masks from the original array and combine them element by element:

arr = np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4])
mask_odd = arr % 2 != 0
mask_positive = arr > 0
result = arr[mask_odd & mask_positive]
print(result)
[1 3]
Mixed Data Types

Why can this array not be compared directly with the integer 2?

arr = np.array(["a", 1, "b", 2, "c", 3, "d", 4])
result = arr[arr > 2]

NumPy chooses one dtype for the complete array. Because strings are present, all values become strings, so comparing them with an integer is invalid. If mixed Python types are required, keep a list and test each element explicitly:

values = ["a", 1, "b", 2, "c", 3, "d", 4]
result = [value for value in values if type(value) is int and value > 2]
print(result)
[3, 4]
Three Ways to Select Positions

Predict all three outputs and explain why the first differs from the other two.

arr = np.array([1, 2, 3, 4, 5])

print(arr[1::2])
print(arr[arr % 2 != 0])
print(arr[(arr.astype(int) / 2) != (arr / 2)])
[2 4]
[1 3 5]
[]

The outputs are [2 4], [1 3 5], and [1 3 5]. The first expression selects odd indices. The other expressions select odd values, although arr % 2 != 0 is much clearer than comparing two division results.

Bonus Points

The professor gives the same bonus to every student, but no grade can go above 10. Write a function named apply_bonus(grades, bonus) that, without loops:

  1. Adds the bonus to every grade.
  2. Uses a mask to cap every grade above 10 at exactly 10.
  3. Returns the new grades and a Boolean array saying which students pass (grade of at least 5).

The original grades array must not be modified.

Sample input: apply_bonus(np.array([9.5, 4.5, 7, 3]), 1)
Expected result: ([10. 5.5 8. 4. ], [ True True True False])

def apply_bonus(grades, bonus):
    new_grades = grades + bonus          # a new array: grades is not modified
    new_grades[new_grades > 10] = 10
    return new_grades, new_grades >= 5


print(apply_bonus(np.array([9.5, 4.5, 7, 3]), 1))
(array([10. ,  5.5,  8. ,  4. ]), array([ True,  True,  True, False]))

grades + bonus creates a new array, so the assignment with the mask only changes the copy.

Character Replacement, Revisited

In Module 1 you replaced every occurrence of the first character of a string with "$", except the first one. Do it again with a mask: turn the string into an array of characters, build a Boolean mask, switch off its first position, and assign "$".

Sample input 1: replace_first("excellent")
Expected result 1: "exc$ll$nt"

Sample input 2: replace_first("crocodile")
Expected result 2: "cro$odile"

def replace_first(s):
    chars = np.array(list(s))
    mask = chars == chars[0]
    mask[0] = False
    chars[mask] = "$"
    return "".join(chars)


print(replace_first("excellent"))
print(replace_first("crocodile"))
exc$ll$nt
cro$odile
Prime Number Checker, Revisited

Write a function is_prime(n) without loops. Build the array of candidate divisors from 2 to \(\sqrt{n}\), compute all the remainders at once, and check that none of them is zero.

Sample input: [is_prime(k) for k in [1, 2, 7, 9, 97, 100]]
Expected result: [False, True, True, False, True, False]

def is_prime(n):
    if n < 2:
        return False
    divisors = np.arange(2, int(n ** 0.5) + 1)
    return bool(np.all(n % divisors != 0))


print([is_prime(k) for k in [1, 2, 7, 9, 97, 100]])
[False, True, True, False, True, False]

For n = 2 and n = 3 the array of divisors is empty, and np.all() of an empty array is True.

Amicable Numbers, Revisited

In Module 1 you wrote sum_of_divisors(num) and are_amicable(num1, num2) with loops. Write them again without loops:

  1. sum_of_divisors(num) builds the array of candidates from 1 to num, uses a mask to keep those that divide num, and returns their sum.
  2. are_amicable(num1, num2) uses it to check whether two different numbers are amicable: the sum of the proper divisors (all divisors except the number itself) of each one equals the other.

Sample input 1: sum_of_divisors(12), sum_of_divisors(16)
Expected result 1: 28, 31

Sample input 2: are_amicable(220, 284), are_amicable(1184, 1210), are_amicable(6, 6)
Expected result 2: True, True, False

def sum_of_divisors(num):
    candidates = np.arange(1, num + 1)
    return np.sum(candidates[num % candidates == 0])


def are_amicable(num1, num2):
    if num1 == num2:
        return False
    return (sum_of_divisors(num1) - num1 == num2) and (sum_of_divisors(num2) - num2 == num1)


print(sum_of_divisors(12), sum_of_divisors(16))
print(are_amicable(220, 284), are_amicable(1184, 1210), are_amicable(6, 6))
28 31
True True False

We subtract the number itself to keep only its proper divisors.

Multiples of 3 or 5

Inspired by Project Euler, Problem 1.

The natural numbers below 10 that are multiples of 3 or 5 are 3, 5, 6, and 9, and their sum is 23.

Write a Python function multiples_3_or_5(n) that returns the sum of all the natural numbers below n that are multiples of 3 or 5. Use a mask instead of a loop.

Sample input 1: multiples_3_or_5(10)
Expected result 1: 23

Sample input 2: multiples_3_or_5(20)
Expected result 2: 78

Sieve of Eratosthenes

The Sieve of Eratosthenes is an ancient method to find all the prime numbers up to a given limit \(n\). Write down all the numbers from 2 to \(n\). Take the first number that is not crossed out, 2, and cross out all its multiples. Then take the next number that is not crossed out, 3, and cross out all its multiples. Repeat until you reach \(\sqrt{n}\): the numbers that remain are the primes.

Write a Python function sieve(n) that returns an array with all the primes up to n, using a Boolean mask to keep track of the crossed-out numbers. You may loop over the numbers up to \(\sqrt{n}\), but cross out all the multiples of each one in a single assignment.

How many primes are there below 1000?

Sample input: sieve(30)
Expected result: [ 2 3 5 7 11 13 17 19 23 29]


Random Numbers and Fancy Indexing

A 50% Event

Complete the condition so "Welcome" is printed with probability 0.5.

p = np.random.uniform(0, 100)

if ...:
    print("Welcome")
p = np.random.uniform(0, 100)

if p >= 50:
    print("Welcome")

A continuous uniform sample is equally likely to fall below or above 50.

Simulating Dice Rolls

Fix the code so each player receives an integer from 1 through 6.

players = ["A", "B", "C"]
results = np.random.uniform(1, 6, size=3)
players = ["A", "B", "C"]
results = np.random.randint(1, 7, size=len(players))
player_results = dict(zip(players, results))
print(player_results)
{'A': np.int64(1), 'B': np.int64(4), 'C': np.int64(6)}

The upper bound 7 is excluded, so the possible results are 1 through 6.

Replacing the First and Last Values

Fix the indexing so both positions become 100 in one assignment.

arr = np.array([1, 2, 3, 4, 5])
arr[0, -1] = 100
arr = np.array([1, 2, 3, 4, 5])
arr[[0, -1]] = 100
print(arr)
[100   2   3   4 100]

The outer brackets create one array of fancy indices for the one-dimensional array.

Repeated Fancy Indices

Predict the output.

arr = np.array([1, 2, 3, 4, 5])
print(arr[[0, 0, 1]])
[1 1 2]

Python prints [1 1 2]. Fancy indices may repeat, and the output preserves their order.

Correcting Row and Column Indices

Fix the code so it extracts the third row and second column.

matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

third_row = matrix[3, :]
second_column = matrix[:, 2]
matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])
third_row = matrix[2, :]
second_column = matrix[:, 1]

print(third_row)
print(second_column)
[7 8 9]
[2 5 8]

Python indices begin at zero.

Extracting a Diagonal

Use fancy indexing to extract the diagonal of this matrix.

matrix = np.array([
    [10, 20, 30, 40],
    [50, 60, 70, 80],
    [90, 100, 110, 120],
    [130, 140, 150, 160]
])
indices = np.arange(matrix.shape[0])
diagonal = matrix[indices, indices]
print(diagonal)
[ 10  60 110 160]

np.diag(matrix) is also valid, but the paired indices make the fancy-indexing mechanism visible.

Reversing an Array

Fix the slice so it reverses this one-dimensional array.

arr = np.array([1, 2, 3, 4, 5])
reversed_arr = arr[::-1, :]
arr = np.array([1, 2, 3, 4, 5])
reversed_arr = arr[::-1]
print(reversed_arr)
[5 4 3 2 1]

A one-dimensional array has only one axis, so a second slice is invalid.

Top-Left Subarray

Extract the top-left 2 × 2 block.

matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])
subarray = matrix[0:2, 0:2]
print(subarray)
[[1 2]
 [4 5]]
Random Seeds and Reassignment

Explain what the seed controls and what value type x has at the end.

x = 10
np.random.seed(x)
x = np.random.randint(low=0, high=1000)
print(x)
265

The seed fixes the pseudo-random sequence, so rerunning the code with the same NumPy random generator produces the same result. The final assignment replaces x = 10 with one NumPy integer sampled from 0 through 999.

Monty Hall Columns

The value 1 marks the winning door in each game. Select door number 2 from every game and compute the overall proportion of winning entries.

games = np.array([
    [0, 1, 0],
    [1, 0, 0],
    [0, 1, 0]
])
door_two = games[:, 1]
gain_rate = np.mean(games)

print(door_two)
print(gain_rate)
[1 0 1]
0.3333333333333333

Door number 2 uses index 1. Because each row contains one car, one third of all entries are winning entries.

Passing by Guessing

The multiple-choice part of an exam has 6 questions with 4 options each. A student who did not study answers every question at random. What is the probability that they get at least half of the questions right?

Write a function named guess_probability(n_students, n_questions, n_options) that, without loops:

  1. Generates a random answer key, one correct option per question.
  2. Generates the random answers of n_students students in a 2D array.
  3. Compares the answers with the key to count the correct answers of each student.
  4. Returns the fraction of students with at least half of the questions right.

Sample input: guess_probability(100_000, 6, 4)
Expected result: approximately 0.17 (it will change slightly with every run)

def guess_probability(n_students, n_questions, n_options):
    key = np.random.randint(0, n_options, size=n_questions)
    answers = np.random.randint(0, n_options, size=(n_students, n_questions))

    correct = np.sum(answers == key, axis=1)   # key is broadcast to every row
    return np.mean(correct >= n_questions / 2)


print(guess_probability(100_000, 6, 4))
0.1681

The mean of a Boolean array is the fraction of True values. The exact probability, from the binomial distribution, is about 0.169.

Encrypted Code, Revisited

In Module 2 you decoded messages with a dictionary of substitution rules. Now build a random substitution cipher with NumPy. Work with lowercase messages without spaces, and represent each letter by its position in the alphabet ("a" is 0, "z" is 25; ord(ch) - ord("a") gives it).

alphabet = np.array(list("abcdefghijklmnopqrstuvwxyz"))
  1. Write make_key(seed) that fixes the random seed and returns a random permutation of the numbers 0 to 25 (np.random.permutation()). The letter at position i is encrypted as the letter at position key[i].
  2. Write encrypt(message, key) that turns the message into an array of positions and encrypts all the letters at once with fancy indexing.
  3. Write decrypt(message, key). Hint: np.argsort(key) gives the inverse permutation.

Sample input:

key = make_key(42)
secret = encrypt("helloworld", key)
print(secret)
print(decrypt(secret, key))

Expected result: blmmekewmy and helloworld

alphabet = np.array(list("abcdefghijklmnopqrstuvwxyz"))


def make_key(seed):
    np.random.seed(seed)
    return np.random.permutation(26)


def to_positions(message):
    return np.array([ord(ch) - ord("a") for ch in message])


def encrypt(message, key):
    return "".join(alphabet[key[to_positions(message)]])


def decrypt(message, key):
    inverse = np.argsort(key)
    return "".join(alphabet[inverse[to_positions(message)]])


key = make_key(42)
secret = encrypt("helloworld", key)
print(secret)
print(decrypt(secret, key))
blmmekewmy
helloworld

key[positions] applies the substitution to every letter at once. If key[i] = j, then np.argsort(key)[j] = i, which undoes the substitution.

Secret Santa

A group of n friends write their names on pieces of paper, shuffle them, and each one draws a name. The draw fails if somebody draws their own name. What is the probability that the draw works?

Write a Python function secret_santa(n, n_trials) that simulates n_trials random draws and returns the fraction of them where nobody draws their own name. Do not use loops. Hint: a draw can be represented as a random permutation of np.arange(n), where friend i gives a present to friend perm[i].

Try it with a growing number of friends. Which number does the probability approach?

Sample input: secret_santa(10, 100_000)
Expected result: approximately 0.37 (it will change slightly with every run)


Statistics and Simulation

Row Means

Fix the code so it prints [2, 5, 8].

arr = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

mean = int(np.mean(arr))
print(mean)
5
row_means = np.mean(arr, axis=1).astype(int)
print(row_means)
[2 5 8]

Reducing axis 1 combines the values across each row.

Means Along Every Axis

Fix the loop so it stores the mean along axis 0 and axis 1 in a dictionary.

arr = np.array([[1, 2], [1, 3], [1, 7]])
axes = arr.shape
dict_means = {}

for axis in axes:
    key = f"mean_{axis}"
    value = np.mean(arr, axis=axis)
    dict_stats.append({key: value})
arr = np.array([[1, 2], [1, 3], [1, 7]])
dict_means = {}

for axis in range(arr.ndim):
    dict_means[f"mean_{axis}"] = np.mean(arr, axis=axis)

print(dict_means)
{'mean_0': array([1., 4.]), 'mean_1': array([1.5, 2. , 4. ])}

arr.shape contains axis sizes (3, 2), not axis numbers. A 2D array has axes 0 and 1.

Normalizing a Vector

The normalized vector should have mean 0. Correct the formula.

arr = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])
mean = np.mean(arr)
std = np.std(arr)
arr_norm = arr - mean / std
print(np.mean(arr_norm))
3.063508326896292
arr_norm = (arr - mean) / std
print(np.mean(arr_norm))
0.0

The subtraction must happen before division. Floating-point arithmetic may display a value extremely close to zero rather than exactly zero.

Accumulation or One Final Total?

Predict each output. Which NumPy expression reproduces the final value of the loop?

arr = np.array([1, 2, 3, 4, 5])
result = 1

for value in arr:
    result = result - value

print(result)
print(1 - np.cumsum(arr))
print(1 - np.sum(arr))
print(1 - np.diff(arr))
-14
[  0  -2  -5  -9 -14]
-14
[0 0 0 0]
  • The loop ends at -14.
  • 1 - np.cumsum(arr) returns [0, -2, -5, -9, -14].
  • 1 - np.sum(arr) returns -14, matching the loop’s final value.
  • 1 - np.diff(arr) returns [0, 0, 0, 0].
Exercise Report

Write a function named exercise_report(marks) that receives a 2D array of marks (one row per student, one column per exercise) and returns three values:

  1. The average mark of each exercise.
  2. The number (starting at 1) of the hardest exercise, the one with the lowest average.
  3. The number (starting at 1) of the exercise whose marks have the highest correlation with the total score of the students. This exercise is the one that best separates strong and weak students.

You may loop over the exercises to compute the correlations. Assume no exercise has the same mark for every student.

Sample input:

marks = np.array([
    [1, 0.5, 1, 0],
    [0.5, 0, 1, 1],
    [1, 1, 1, 0],
    [0, 0, 1, 0.5],
    [1, 0.5, 0, 1],
])
exercise_report(marks)

Expected result: ([0.7 0.4 0.8 0.5], 2, 1)

def exercise_report(marks):
    averages = np.mean(marks, axis=0)
    hardest = np.argmin(averages) + 1

    total = np.sum(marks, axis=1)
    correlations = np.zeros(marks.shape[1])
    for j in range(marks.shape[1]):
        correlations[j] = np.corrcoef(marks[:, j], total)[0, 1]
    best = np.argmax(correlations) + 1

    return averages, hardest, best


marks = np.array([
    [1, 0.5, 1, 0],
    [0.5, 0, 1, 1],
    [1, 1, 1, 0],
    [0, 0, 1, 0.5],
    [1, 0.5, 0, 1],
])
print(exercise_report(marks))
(array([0.7, 0.4, 0.8, 0.5]), np.int64(2), np.int64(1))
Bus Lines, Revisited

In Module 2 you sorted bus lines by how far their route time is from the average time. Do it again with NumPy, without loops. Write sort_lines(names, times) that:

  1. Computes the average route time and the distance of each line from it.
  2. Sorts the lines by that distance, from closest to farthest.
  3. When two lines are equally distant, places the faster one first. Hint: np.lexsort((times, distances)) sorts by distances and breaks ties with times.
  4. Returns the sorted names and times as two arrays.

Sample input 1: sort_lines(["Line 1", "Line 2", "Line 3"], [50, 30, 40])
Expected result 1: ['Line 3' 'Line 2' 'Line 1'] and [40 30 50]

Sample input 2: sort_lines(["A", "B", "C", "D", "E"], [20, 35, 50, 30, 40])
Expected result 2: ['B' 'D' 'E' 'A' 'C'] and [35 30 40 20 50]

def sort_lines(names, times):
    names = np.array(names)
    times = np.array(times)
    distances = np.abs(times - np.mean(times))

    order = np.lexsort((times, distances))  # last key is the primary key
    return names[order], times[order]


print(sort_lines(["Line 1", "Line 2", "Line 3"], [50, 30, 40]))
print(sort_lines(["A", "B", "C", "D", "E"], [20, 35, 50, 30, 40]))
(array(['Line 3', 'Line 2', 'Line 1'], dtype='<U6'), array([40, 30, 50]))
(array(['B', 'D', 'E', 'A', 'C'], dtype='<U1'), array([35, 30, 40, 20, 50]))
Dice Duel

Inspired by Project Euler, Problem 205.

Alice rolls five four-sided dice (faces 1 to 4) and Bob rolls three six-sided dice (faces 1 to 6). The player with the highest total wins, and equal totals are a tie.

Write a Python function dice_duel(n_games) that simulates n_games games without loops, and returns the estimated probabilities that Alice wins, that Bob wins, and that they tie. How much do the estimates change when you simulate 100, 1,000, or 100,000 games?

Sample input: dice_duel(100_000)
Expected result: approximately (0.65, 0.26, 0.09) (it will change slightly with every run)


Linear Algebra

A Matrix Times Its Inverse

Predict z. Why might some entries not display as exact integers?

x = np.array([[1, 2], [3, 4]])
y = np.linalg.inv(x)
z = x @ y
print(z)
[[1.0000000e+00 0.0000000e+00]
 [8.8817842e-16 1.0000000e+00]]

Mathematically, z is the identity matrix. Floating-point calculations may produce tiny errors such as 2.22044605e-16 instead of zero. Use np.allclose(z, np.eye(2)) to test approximate equality:

print(np.allclose(z, np.eye(2)))
True

Converting directly to integers is unsafe because conversion truncates rather than rounds.

Finding Eigenvalues

Which call returns the eigenvalues and eigenvectors of a square matrix?

X = np.array([[1, 2], [3, 4]])

Choose among np.eigen(X), np.eigvals(X), and np.linalg.eig(X).

eigenvalues, eigenvectors = np.linalg.eig(X)
print(eigenvalues)
print(eigenvectors)
[-0.37228132+0.j  5.37228132+0.j]
[[-0.82456484+0.j -0.41597356+0.j]
 [ 0.56576746+0.j -0.90937671+0.j]]

np.linalg.eig() returns both results. np.linalg.eigvals() returns only eigenvalues; np.eigvals() and np.eigen() do not exist.

Recover the Weights

A colleague published the total points of three students, but forgot to say how many points each of the three exercises was worth. You know the marks (between 0 and 1) of each student. Write a function named recover_weights(marks, scores) that finds the points of each exercise by solving a system of linear equations with np.linalg.solve().

Sample input:

marks = np.array([[1, 0.5, 0], [0.5, 1, 1], [1, 1, 0.5]])
scores = np.array([65, 75, 90])
recover_weights(marks, scores)

Expected result: [50. 30. 20.]

Explain why this would not work if two students had exactly the same marks.

def recover_weights(marks, scores):
    return np.linalg.solve(marks, scores)


marks = np.array([[1, 0.5, 0], [0.5, 1, 1], [1, 1, 0.5]])
scores = np.array([65, 75, 90])
print(recover_weights(marks, scores))
[50. 30. 20.]

If two students had the same marks, two rows of the matrix would be equal. The matrix would be singular (determinant zero), the system would have fewer independent equations than unknowns, and np.linalg.solve() would raise a LinAlgError.

Fibonacci with Matrices

Inspired by Project Euler, Problem 2.

The Fibonacci numbers can be computed with powers of a matrix:

\[\begin{pmatrix} 1 & 1 \\ 1 & 0 \end{pmatrix}^n = \begin{pmatrix} F_{n+1} & F_n \\ F_n & F_{n-1} \end{pmatrix}\]

Write a Python function fibonacci(n) that returns \(F_n\) using np.linalg.matrix_power(). Use it to add up the even Fibonacci numbers that do not exceed a given limit.

Then compute the eigenvalues of the matrix. Do you recognize them? Binet’s formula, \(F_n = \frac{\varphi^n - \psi^n}{\sqrt{5}}\), uses these two eigenvalues to compute \(F_n\) directly. Compare it with your function for \(n = 1, \dots, 80\). Do they always agree?

Sample input 1: fibonacci(10)
Expected result 1: 55

Sample input 2: the even Fibonacci numbers not exceeding 100
Expected result 2: 44 (that is, \(2 + 8 + 34\))


Vectorization

Element-Wise Comparison

Predict c.

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])
c = a > b

c is [False, False, False]. NumPy performs the comparison element by element.

Broadcasting a Column and a Row

Predict the shape and values of c.

a = np.array([[1], [2], [3]])
b = np.array([4, 5, 6])
c = a > b

a has shape (3, 1) and b has shape (3,), which behaves like (1, 3). Broadcasting produces a (3, 3) array. Every comparison is false because each value in a is smaller than every value in b.

Broadcasting Failure

Why does this comparison raise a ValueError?

a = np.array([1, 2, 3])
b = np.array([4, 5, 6, 7])
c = a > b

The one-dimensional shapes (3,) and (4,) are incompatible because their only dimensions differ and neither equals 1.

Vectorize the Grader

The following function computes the final grade of each student, applying a 10% penalty to those who delivered late. Write a function named grade_vectorized(marks, weights, late) that returns the same result without any loop, where marks is a 2D array, weights a 1D array, and late a Boolean array with one value per student.

def grade_loop(marks, weights, late):
    grades = []
    for i in range(len(marks)):
        total = 0
        for j in range(len(weights)):
            total = total + marks[i][j] * weights[j]
        grade = total / 10
        if late[i]:
            grade = grade * 0.9
        grades.append(round(grade, 2))
    return grades

Sample input:

marks = np.array([[1, 0.5, 1, 0.5], [0.75, 1, 0, 1], [0.5, 0.5, 0.5, 0.5]])
weights = np.array([40, 30, 20, 10])
late = np.array([False, True, False])
grade_vectorized(marks, weights, late)

Expected result: [8. 6.3 5. ]

Then create random marks for 10,000 students and 11 exercises, check that both functions return the same grades, and compare their running times with %timeit.

def grade_vectorized(marks, weights, late):
    grades = marks @ weights / 10
    grades[late] = grades[late] * 0.9
    return np.round(grades, 2)


marks = np.array([[1, 0.5, 1, 0.5], [0.75, 1, 0, 1], [0.5, 0.5, 0.5, 0.5]])
weights = np.array([40, 30, 20, 10])
late = np.array([False, True, False])
print(grade_vectorized(marks, weights, late))
[8.  6.3 5. ]

To compare with the loop version on a large class:

big_marks = np.random.rand(10_000, 11)
big_weights = np.array([20, 20, 10, 10, 10, 5, 5, 5, 5, 5, 5])
big_late = np.random.rand(10_000) < 0.1

print(np.allclose(grade_loop(big_marks, big_weights, big_late),
                  grade_vectorized(big_marks, big_weights, big_late)))

%timeit grade_loop(big_marks, big_weights, big_late)
%timeit grade_vectorized(big_marks, big_weights, big_late)
Pythagorean Triples, Revisited

In Module 1 you checked whether three given numbers form a Pythagorean triple. Now find all of them at once, without loops. Write pythagorean_triples(n) that returns a 2D array with one row [a, b, c] for every triple with \(a < b\) and \(c \le n\):

  1. Build a column of values a and a row of values b, both from 1 to n, and let broadcasting create the table of all \(a^2 + b^2\).
  2. Compute the candidate c with np.sqrt() and keep only the cells where it is an integer, \(a < b\), and \(c \le n\).
  3. Stack the selected a, b, and c as the three columns of the result.

Sample input: pythagorean_triples(20)
Expected result:

[[ 3  4  5]
 [ 5 12 13]
 [ 6  8 10]
 [ 8 15 17]
 [ 9 12 15]
 [12 16 20]]
def pythagorean_triples(n):
    a = np.arange(1, n + 1).reshape(-1, 1)   # column
    b = np.arange(1, n + 1).reshape(1, -1)   # row
    c_squared = a**2 + b**2                  # table (n, n)
    c = np.round(np.sqrt(c_squared)).astype(int)

    mask = (c**2 == c_squared) & (a < b) & (c <= n)
    rows, cols = np.nonzero(mask)
    return np.column_stack((rows + 1, cols + 1, c[mask]))


print(pythagorean_triples(20))
[[ 3  4  5]
 [ 5 12 13]
 [ 6  8 10]
 [ 8 15 17]
 [ 9 12 15]
 [12 16 20]]

np.nonzero(mask) returns the row and column indices of the True cells; adding 1 turns indices into the values of a and b.

Low Points

Inspired by Advent of Code 2021, Day 9.

A height map of a cave floor is stored as a 2D array of digits, from 0 (lowest) to 9 (highest). A low point is a cell that is strictly lower than all its neighbours: up, down, left, and right. Diagonals do not count, and cells on the border have fewer neighbours.

Write a Python function low_points(heights) that returns a Boolean mask of the low points, without loops. Hint: np.pad(heights, 1, constant_values=10) surrounds the map with a border of 10s, so every cell has four neighbours.

The risk level of a low point is its height plus one. What is the sum of the risk levels of all the low points?

Sample input:

heights = np.array([
    [5, 3, 4, 8, 9, 7, 6, 5, 6, 7],
    [4, 6, 7, 8, 9, 8, 5, 4, 5, 6],
    [6, 7, 8, 9, 7, 6, 7, 8, 9, 3],
    [8, 8, 9, 6, 5, 6, 8, 9, 8, 4],
    [9, 7, 6, 5, 6, 7, 9, 9, 7, 5],
])

Expected result: 6 low points with heights [3 4 4 3 5 5], and a total risk level of 30.


Further Practice