NumPy arrays and vectorized operations: Exercises

Before we start, three tiny bits of vocabulary that the exercises reuse. Read these once; each is one sentence.

Definition Three words we will keep using

ufunc ("universal function"): a NumPy function that runs element-by-element over a whole array in compiled C, and broadcasts its inputs. Examples: np.exp, np.sin, np.maximum, and even + and ** (they are ufuncs under the hood). The parent note NumPy arrays and vectorized operations introduces these as "vectorized operations" — same thing.
None as a new axis: writing None (its formal name is np.newaxis) inside an index inserts a length-1 dimension at that spot. So if v has shape (3,), then v[:, None] has shape (3, 1) and v[None, :] has shape (1, 3). We use this to line arrays up on different axes so broadcasting can pair them.
empty / zero-length dimension: a shape can legally contain a 0, e.g. (0, 3) is "zero rows, each of width 3" — an array holding no elements but still knowing its width. It shows up when a filter matches nothing.

Now, one picture fixes the single rule that half of these problems test.

The One Rule: Broadcasting, drawn

Figure — NumPy arrays and vectorized operations

How to read the figure (it drives Exercises 3.1 and 4.1). The two shapes (4,1,3) and (5,3) are stacked with their right edges flush — this is step 1 of the rule made visual. (5,3) is shorter, so a 1 is prepended (step 2) to make (1,5,3). Now walk the three columns:

Rightmost column 3 vs 3: equal, no stretch needed → keep 3.
Middle column 1 vs 5: the gray arrow shows array A's length-1 axis being virtually repeated 5 times (step 4) → becomes 5.
Leftmost column 4 vs 1: array B's length-1 axis is repeated 4 times → becomes 4.

The green row at the bottom is the output shape (4,5,3), and each green cell is literally the max of the two cells above it. The red note reminds you of the only failure mode: a column with two different numbers where neither is 1 — no legal repetition exists. Everything below refers back to this exact walk.

Keep this picture in your head; problems L1–L5 all lean on it.

L1 — Recognition

Exercise 1.1 (L1)

State the shape and dtype NumPy infers for:

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

Recall Solution

Two rows, three columns, so ==shape = (2, 3). Every entry is a whole number, so NumPy picks an integer type: dtype = int64== (on 64-bit platforms). Why: shape is read outer-list-length first (2 sub-lists), then inner-list-length (3 numbers each). dtype is the smallest common type that holds every element without loss — all ints ⇒ int64.

Exercise 1.2 (L1)

What does each line return?

np.arange(0, 10, 2)      # (a)
np.linspace(0, 1, 5)     # (b)
np.zeros((2, 3)).dtype   # (c)

Recall Solution

(a) array([0, 2, 4, 6, 8]) — start at 0, step 2, stop before 10.
(b) array([0. , 0.25, 0.5 , 0.75, 1. ]) — 5 points, endpoints included, gaps of $\frac{1 - 0}{5 - 1} = 0.25$ .
(c) float64 — zeros/ones default to floating point. Why (b) has step 0.25 and not 0.2: linspace(a, b, N) places N points, so there are N-1 gaps: $\frac{b - a}{N - 1}$ .

Exercise 1.3 (L1)

For A of shape (3, 4), what shape does the slice A[1:3, :] have, and is it a view or a copy?

Recall Solution

Rows 1 and 2 (Python stops before 3), all columns ⇒ shape ==(2, 4). It is a view== — it shares the parent's memory buffer. Writing into it writes into A. Why a view: basic slicing only changes the start offset and strides, never copies data. This is why slicing multi-GB arrays is instant.

L2 — Application

Exercise 2.1 (L2)

Given

a = np.array([1, 2, 3])       # shape (3,)
b = np.array([[10], [20]])    # shape (2, 1)
c = a + b

Write out c fully and give its shape.

Recall Solution

First write each shape explicitly. a is a 1-D array of length 3, so its shape is (3,). b is 2-D, shape (2, 1). Align from the right and prepend a 1 to the shorter shape so both have the same length: $a : (3,) ⟶ (1, 3), b : (2, 1)$ Now stack them, right edges flush, and read each column: $b : a : max : col 0 212 col 1 133$

Column 0: 2 vs 1 → the 1 (from a) stretches up to 2 rows.
Column 1: 1 vs 3 → the 1 (from b) stretches across to 3 columns. So (2, 1) vs (1, 3) → output shape ==(2, 3)==. Row 0 uses b=10: [11, 12, 13]. Row 1 uses b=20: [21, 22, 23]. $c = [112112221323]$ Why: a (a single row) is repeated down to fill 2 rows; b (a single column) is repeated across to fill 3 columns. Every cell is a[j] + b[i] — an outer sum.

Exercise 2.2 (L2)

Replace this loop with one vectorized line, and give the result for data = np.array([-3, -1, 0, 2, 5]):

out = np.empty_like(data)
for i in range(len(data)):
    out[i] = data[i] if data[i] > 0 else 0

Recall Solution

This is the ReLU function. Vectorized:

out = np.maximum(data, 0)

Result: ==array([0, 0, 0, 2, 5])==. Why np.maximum and not max: np.maximum is a ufunc (defined at the top of this page — a compiled, element-wise function). It broadcasts the scalar 0 against every element and runs the whole comparison in C. Plain max(data, 0) tries to compare the whole array object to 0 and raises an error, and a Python loop would be ~100× slower for the same result.

Exercise 2.3 (L2)

Given x = np.array([10, 20, 30, 40, 50]), use boolean indexing to return only the elements strictly greater than 25, then their sum.

Recall Solution

mask = x > 25          # array([False, False, True, True, True])
x[mask]                # array([30, 40, 50])
x[mask].sum()          # 120

Selected elements: ==[30, 40, 50], sum 120==. Why: x > 25 builds a boolean array of the same shape; x[mask] keeps positions where the mask is True. This creates a copy (the survivors are not contiguous in memory).

L3 — Analysis

Exercise 3.1 (L3)

Do these broadcast? If yes, give the output shape; if no, say which aligned column breaks the rule.

(a) (3, 4)    with (4,)
(b) (3, 4)    with (5,)
(c) (4, 1, 3) with (5, 3)
(d) (2, 3)    with (2, 1)
(e) (0, 3)    with (3,)      # zero-length first dimension

Recall Solution

(a) align: (3,4) vs (1,4) → columns (3,1) ok, (4,4) ok ⇒ (3, 4).
(b) align: (3,4) vs (1,5) → last column (4,5): neither equals nor is 1 ⇒ FAILS (the 4 vs 5 column).
(c) align: (4,1,3) vs (1,5,3) → columns (4,1)→4, (1,5)→5, (3,3)→3 ⇒ (4, 5, 3) (this is exactly the figure at the top).
(d) columns (2,2) ok, (3,1) ok ⇒ (2, 3).
(e) align: (0,3) vs (1,3) → column 0 is 0 vs 1: one side is 1, so it is compatible and stretches to max(0,1)=0; column 1 is 3 vs 3 ok ⇒ (0, 3). The result is a valid empty array — no error, just zero rows. Why (b) breaks: two "real" sizes 4 and 5 can only match if identical; there is no size-1 dimension to stretch, so no legal repetition exists (the red column in the figure). Why (e) does NOT break: the rule is equal, or one is 1. Here 0 vs 1 satisfies "one is 1", and the output takes max(0,1)=0. A 0 only conflicts with a size ≥ 2 — e.g. (0,3) vs (2,3) would fail on the first column (0 vs 2, neither is 1).

Exercise 3.2 (L3)

v = np.array([1., 2., 3., 4.]). You want to normalize it to unit length. Explain why

v / v.sum()

is wrong and write the correct expression. Give the numeric result of the correct one.

Recall Solution

v.sum() = 10, so v / v.sum() = [0.1, 0.2, 0.3, 0.4] — those sum to 1, but that is probability normalization, not length normalization. "Unit length" means the Euclidean norm equals 1: $∥ v ∥ = \sum_{i} v_{i}^{2} = 1 + 4 + 9 + 16 = 30 \approx 5.4772$ Correct:

v / np.linalg.norm(v)

Result: [0.1826, 0.3651, 0.5477, 0.7303] (each is $v_{i} / 30$ ). Check: their squares sum to 1. Why the tool: np.linalg.norm computes $\sum v_{i}^{2}$ , the geometric length; dividing by it rescales the vector to length 1 while keeping its direction.

Exercise 3.3 (L3)

A has shape (1000, 3) (1000 points in 3-D). You want to subtract the column mean from each column (center the data). Why does A - A.mean(axis=0) work but A - A.mean(axis=1) fail, and what shapes appear?

Recall Solution

A.mean(axis=0) collapses the rows (axis 0), leaving one mean per column: shape (3,). Broadcasting (1000,3) vs (3,) → align (1000,3) vs (1,3) ⇒ works, output (1000, 3).
A.mean(axis=1) collapses the columns, giving one mean per row: shape (1000,). Broadcasting (1000,3) vs (1000,) → align (1000,3) vs (1,1000) → last column (3, 1000) ⇒ FAILS. Fix for row-centering (if you truly wanted it): keep the axis with keepdims=True so the shape is (1000, 1), which does broadcast against (1000, 3). Why axis=0: "axis 0" is the axis you move along and destroy. Averaging along axis 0 walks down the 1000 rows and returns a value for each of the 3 columns — exactly the per-feature mean ML wants.

L4 — Synthesis

Exercise 4.1 (L4)

Compute the pairwise squared Euclidean distance between every pair of the points

P = np.array([[0., 0.],
              [3., 0.],
              [0., 4.]])   # shape (3, 2): three 2-D points

using only broadcasting (no Python loop). Give the full 3×3 result matrix D where D[i,j] = ||P[i] - P[j]||².

Recall Solution

The trick: make the row-index and column-index live on different axes, then let broadcasting form all pairs. We use the None-new-axis tool defined at the top of the page.

diff = P[:, None, :] - P[None, :, :]   # shape (3,1,2) - (1,3,2) -> (3,3,2)
D = (diff ** 2).sum(axis=2)            # sum over the coordinate axis -> (3,3)

Why the Nones. P has shape (3,2). Writing P[:, None, :] inserts a length-1 axis in the middle, giving (3,1,2) — "points down axis 0". Writing P[None, :, :] inserts a length-1 axis at the front, giving (1,3,2) — "points across axis 1". Broadcasting stretches both length-1 axes to size 3, producing (3,3,2), so diff[i,j] is exactly the vector $P_{i} - P_{j}$ . Squaring and summing the last axis (the two coordinates) gives the squared distance. Distances: $P_{0}$ – $P_{1}$ : $3^{2} = 9$ ; $P_{0}$ – $P_{2}$ : $4^{2} = 16$ ; $P_{1}$ – $P_{2}$ : $3^{2} + 4^{2} = 25$ . $D = 0916902516250$ Why this beats a double loop: the loop is $O (n^{2})$ in Python; the broadcast version is $O (n^{2})$ in C — same math, ~100× faster, and it is the exact pattern behind k-nearest-neighbours and attention. See 2.1.03-Matrix-operations-for-ML.

How to read the figure (it is the geometry behind the matrix D). The three coloured dots are the rows of P plotted on the plane. Each gray line is one difference vector $P_{i} - P_{j}$ that the broadcast diff array holds; the red label on each line is its squared length, which is precisely the matrix entry D[i,j]. Trace them back to the derivation: the line from $P_{1}$ to $P_{2}$ has horizontal drop 3 and vertical rise 4, so its squared length is $3^{2} + 4^{2} = 25$ — the 25 in D. The diagonal of D is zero (the arrow from a point to itself has length 0), and D is symmetric because the arrow $P_{i} \to P_{j}$ has the same length as $P_{j} \to P_{i}$ .

Exercise 4.2 (L4)

Implement a numerically-safe softmax of a row vector z in vectorized NumPy, and evaluate it for z = np.array([1., 2., 3.]). The softmax of entry $i$ is $σ (z)_{i} = \frac{e ^{z_{i}}}{\sum _{j} e ^{z_{j}}} .$

Recall Solution

z_shift = z - z.max()          # subtract max for numerical stability
e = np.exp(z_shift)            # np.exp is a ufunc: element-wise
soft = e / e.sum()

With z.max() = 3: shifted [-2, -1, 0], exps [0.1353, 0.3679, 1.0], sum 1.5032. Result: ==[0.0900, 0.2447, 0.6652]== (sums to 1). Why subtract the max: exp overflows for large inputs. Subtracting a constant $c$ from every $z_{i}$ multiplies numerator and denominator by $e^{- c}$ , which cancels — the softmax is unchanged, but now the largest exponent is $e^{0} = 1$ , so no overflow. This is a required trick in every deep-learning framework; see 3.1.02-Tensor-operations-in-PyTorch.

L5 — Mastery

Exercise 5.1 (L5)

You must solve $A x = b$ with

A = np.array([[3., 1.], [1., 2.]])
b = np.array([9., 8.])

(1) Solve it by hand and confirm. (2) Explain, in cost and stability terms, why np.linalg.solve(A, b) is preferred over np.linalg.inv(A) @ b.

Recall Solution

(1) By hand. Write the two equations the matrix encodes:

Row 1: $3 x_{1} + x_{2} = 9$
Row 2: $x_{1} + 2 x_{2} = 8$

From row 2, isolate $x_{1}$ : $x_{1} = 8 - 2 x_{2}$ . Substitute into row 1: $3 (8 - 2 x_{2}) + x_{2} = 9 \Rightarrow 24 - 6 x_{2} + x_{2} = 9 \Rightarrow 24 - 5 x_{2} = 9 \Rightarrow x_{2} = 3.$ Back-substitute: $x_{1} = 8 - 2 (3) = 2$ . So $x = [23], check A x = [3 (2) + 1 (3) 1 (2) + 2 (3)] = [98] ✓$ In NumPy:

np.linalg.solve(A, b)   # array([2., 3.])

(2) Cost. For an $n \times n$ system, np.linalg.solve does one LU factorization ( $A = LU$ , with $L$ lower-triangular and $U$ upper-triangular) costing about $\frac{1}{3} n^{3}$ floating-point operations, then a forward substitution (solve $L y = b$ ) and a back substitution (solve $U x = y$ ), each only $O (n^{2})$ . Total $\approx \frac{1}{3} n^{3}$ . By contrast, np.linalg.inv(A) computes the full inverse, which is like solving $n$ separate systems (one for each column of the identity), costing about $n^{3}$ operations — roughly 3× more work — and then you still pay an extra $O (n^{2})$ for the matrix–vector product inv(A) @ b.

Stability. Forming $A^{- 1}$ explicitly and multiplying by it accumulates extra rounding error; the accuracy of any linear solve already degrades with the condition number of $A$ (how much it amplifies input error), and routing through the explicit inverse adds roundoff that the direct substitution in solve avoids. For ill-conditioned $A$ , inv(A) @ b is visibly less accurate than solve(A, b).

Rule of thumb: never invert to solve — factor and substitute. Use np.linalg.solve. This same principle underlies the linear-algebra kernels in 2.1.03-Matrix-operations-for-ML.

Exercise 5.2 (L5)

Rank these three implementations of y = x² for x = np.arange(1_000_000, dtype=float) from fastest to slowest, and justify with the mechanism:

# (A) [xi**2 for xi in x]
# (B) x ** 2
# (C) np.array([xi**2 for xi in x])

Recall Solution

Fastest → slowest: ==(B) < (A) < (C)== (B is fastest).

(B) x ** 2 is a single ufunc call: the whole million-element loop runs in compiled C with SIMD, no Python-level per-element overhead. ~milliseconds.
(A) the list comprehension runs a Python-level loop (interpreter dispatch, boxing each float) — ~100× slower than (B). It also returns a list, not an array.
(C) does everything (A) does and then copies the million-element list into a fresh ndarray — strictly more work than (A), so it is slowest. Why (B) wins the mechanism argument: the cost of a Python loop is dominated by per-iteration interpreter overhead (type checks, refcounts), not by the multiply itself. Vectorization pays that overhead once. This is the whole reason NumPy exists — see 1.4.02-List-comprehensions-and-generators for the loop it replaces.

Exercise 5.3 (L5)

a = np.arange(5). Predict the value of a after each line, and explain the memory difference between lines 1 and 2.

b = a                    # line 1
b[0] = 99                # line 2
c = a.copy()             # line 3
c[1] = -1                # line 4

Recall Solution

After line 2: a = [99, 1, 2, 3, 4]. b = a binds a second name to the same array object — no copy, so mutating b mutates a.
After line 4: a is unchanged at [99, 1, 2, 3, 4]; c = [99, -1, 2, 3, 4]. a.copy() allocates a new buffer, so writing to c cannot touch a. Why it matters: slices and plain assignment are views/aliases (cheap, shared memory); .copy() is an explicit new buffer. Silent aliasing bugs — mutating a "view" you thought was independent — are the #1 source of NumPy correctness errors.

Quick self-test

Cover the right side.

Broadcasting compatibility of (6, 1) and (1, 4)

works → output (6, 4)

Broadcasting compatibility of (3,) and (4,)

fails → 3 vs 4, neither is 1

Broadcasting compatibility of (0, 3) and (3,)

works → output (0, 3), a valid empty array (0 vs 1 → one is 1)

What does v[:, None] do to shape (3,)

makes it (3, 1) — inserts a length-1 axis (that is np.newaxis)

np.maximum(x, 0) implements which activation

ReLU

A ufunc is

a compiled, element-wise NumPy function that broadcasts its inputs (e.g. np.exp, +, **)

Divide a vector by which quantity to get unit length

np.linalg.norm(v)

axis=0 on a (1000, 3) array averages over which entities

the 1000 rows, giving one mean per column

Why subtract z.max() inside softmax

to prevent exp overflow without changing the result

Why solve(A,b) over inv(A)@b

~3× cheaper (one LU vs full inverse) and numerically more stable

b = a; b[0] = 99 — is a changed

yes, b is an alias, not a copy