NumPy for Machine Learning

NumPy provides the ndarray — a fast, fixed-type array — plus vectorized math operations that avoid slow Python loops. Almost every ML library (Pandas, scikit-learn, PyTorch) is built on top of it.

Why Not Just Use Python Lists?

import numpy as np
import time

n = 1_000_000
py_list = list(range(n))
np_array = np.arange(n)

start = time.time()
py_result = [x * 2 for x in py_list]          # pure Python loop
print("Python list:", time.time() - start)

start = time.time()
np_result = np_array * 2                       # vectorized NumPy operation
print("NumPy array:", time.time() - start)

Expected output: the NumPy version typically runs 10–50x faster — it operates on contiguous, fixed-type memory in optimized C code, instead of looping through Python objects one at a time.

Vectorization: The Core Idea

Vectorization means expressing an operation on an entire array at once, instead of writing an explicit loop. This isn't just a performance trick — it's the mental model nearly all ML code is written in.

a = np.array([1, 2, 3, 4])
b = np.array([10, 20, 30, 40])

print(a + b)        # [11 22 33 44] — element-wise, no loop needed
print(a * b)        # [10 40 90 160]
print(a.mean())      # 2.5 — aggregation
print((a > 2))       # [False False True True] — boolean mask

Broadcasting

Broadcasting lets NumPy apply an operation between arrays of different shapes, by conceptually "stretching" the smaller one — without actually copying data.

features = np.array([[100, 2], [200, 3], [150, 5]])   # 3 rows, 2 columns
scale    = np.array([0.01, 1])                          # 1 row, 2 columns — broadcasts to every row

scaled = features * scale
print(scaled)
# [[1.  2.]
#  [2.  3.]
#  [1.5 5.]]

This is exactly the mechanism behind feature scaling operations in Feature Scaling.

Axis — The Most Confusing Beginner Concept

matrix = np.array([[1, 2, 3], [4, 5, 6]])   # 2 rows, 3 columns

print(matrix.sum(axis=0))   # [5 7 9]  — sum DOWN each column
print(matrix.sum(axis=1))   # [6 15]   — sum ACROSS each row

axis=0 collapses rows (operates column-wise); axis=1 collapses columns (operates row-wise). Getting this backwards is one of the most common NumPy/Pandas bugs.

Practical Use Cases

• Storing feature matrices before feeding them to scikit-learn models
• Implementing custom loss functions or gradient calculations from scratch
• Fast statistical summaries (mean, std, correlation) on large numeric arrays

Common Mistakes

• Mixing up axis=0 and axis=1 — always sanity-check with a small example, not just intuition.
• Using Python's for loops over NumPy arrays instead of vectorized operations — defeats the entire purpose of using NumPy.
• Forgetting that slicing a NumPy array returns a view, not a copy — modifying the slice can silently modify the original array. Use .copy() when you need an independent array.

Interview Relevance

Q: "What does vectorization mean, and why does it matter for ML?" It means expressing operations over whole arrays instead of looping element-by-element in Python — it's dramatically faster because the loop runs in optimized C, not the Python interpreter, which matters when training on large datasets.

Practice Question

Given a 2D NumPy array of shape (100, 5) representing 100 samples with 5 features, write the one-line expression to compute the mean of each feature (column) across all samples.

Related ML Notes

• NumPy Arrays
• Pandas for ML
• Feature Scaling