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.

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
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
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 column1/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
Quick check
X has shape (500, 4). What does X - X.mean(axis=0) do?
Go deeper
- 1Practice array creation, indexing, slicing, broadcasting and reshaping.
- 2Rewrite a loop-based computation as a vectorized NumPy operation and compare speed.
- 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.