Stop 03 · Data Manipulation & VisualizationTopics · 1/10

NumPy Arrays & Vectorization

Part of the AI Engineer Roadmap

In short

A Python loop over a million numbers is slow. The same maths on a NumPy array runs many times faster, in one line. pandas, scikit-learn and PyTorch all build on this idea.

NumPy's array type and vectorized operations let you run math on entire datasets without slow Python loops — the performance foundation every ML library is built on.

The purple blob stirs a giant kadhai of biryani while a tiny omelette pan sits to the side.

Think of it like this

The wedding caterer's giant kadhai

Making 500 omelettes one by one would take all night. The caterer cooks biryani for 500 in one huge kadhai, one stir for everyone. A Python loop is the omelette pan. A NumPy operation is the kadhai: one instruction, applied to every value at once.

Watch it happen

Adding 18% GST to every price

100
0
250
1
80
2
40
3

1/4prices is a NumPy array of four numbers, all the same type, packed together in memory.

The big ideas

What you need to know

One type, packed tight

A NumPy array holds values of one type in a single block of memory. That's why it's fast and compact.

Vectorize, don't loop

Write arr * 2, not a for loop. The work happens in compiled C code instead of slow Python steps.

Broadcasting stretches shapes

A (1000, 3) matrix minus a (3,) row subtracts that row from every row. No copying by hand.

Boolean masks filter

arr[arr > 0] keeps only positive values. Masks replace most if-inside-a-loop code.

See it in code

Walk through the code

python
1import time2import numpy as np3 4prices = np.array([100, 250, 80, 40])5print(prices * 1.18)              # [118.  295.   94.4  47.2]6print(prices[prices > 90])        # [100 250]7 8big = np.random.default_rng(0).random(1_000_000)9t = time.perf_counter(); s1 = sum(x * 2 for x in big); loop_s = time.perf_counter() - t10t = time.perf_counter(); s2 = (big * 2).sum(); vec_s = time.perf_counter() - t11print(f"loop {loop_s:.3f}s vs numpy {vec_s:.4f}s")12 13X = np.array([[1, 200], [3, 400], [5, 600]])14print(X - X.mean(axis=0))         # broadcasting: centre each column

1/3Maths and filtering on the whole array with no loop.

Run this in Colab and compare the two timings yourself. The exact numbers depend on your machine, but the gap is always large.

Save yourself the pain

Common mistakes

Oops

Looping over a NumPy array with for and appending to a list.

Do this instead

Express the maths on the whole array. If you need a condition, use a mask or np.where.

Oops

Building an array by calling np.append in a loop.

Do this instead

Collect values in a Python list, then call np.array once. np.append copies the whole array every time.

Your mission

Race the loop

~15 min

Quick check

X has shape (500, 4). What does X - X.mean(axis=0) do?

Go deeper

  1. 1Practice array creation, indexing, slicing, broadcasting and reshaping.
  2. 2Rewrite a loop-based computation as a vectorized NumPy operation and compare speed.
  3. 3Learn broadcasting rules carefully — they're a frequent source of subtle bugs.

More topics at this stop

Stuck on this topic? Ask an Insider

Get 1:1 guidance from people who've walked this exact path — on the InsideEdge app.

Download
InsideEdge