NumPy Basics: Arrays, Vectorisation and Linear Algebra

By Engr. Henry Akpan

Fast Maths on Whole Arrays

Python lists are flexible, but they are slow for number crunching: adding a million pairs of numbers means a million separate Python operations. NumPy (Numerical Python) stores numbers in compact arrays and performs operations on whole arrays at once, in fast compiled code. It is the foundation of scientific Python: pandas, SciPy, Matplotlib, scikit-learn and most machine-learning libraries are built on it. For science and engineering students, NumPy turns Python into a powerful calculator for experiments, statistics and linear algebra.

This tutorial covers NumPy basics step by step: creating arrays, shape and data types, vectorised arithmetic, indexing and slicing, Boolean masks, statistics along axes, broadcasting, reshaping, linear algebra (including solving simultaneous equations), and random numbers. Install it with pip install numpy.

Creating Arrays

import numpy as np

scores = np.array([72, 58, 90, 45, 81])
print(scores, type(scores).__name__, scores.dtype)

print(np.zeros(3))
print(np.ones((2, 3)))
print(np.arange(0, 10, 2))         # like range(), but gives an array
print(np.linspace(0, 1, 5))        # 5 evenly spaced values from 0 to 1

Output:

[72 58 90 45 81] ndarray int64
[0. 0. 0.]
[[1. 1. 1.]
 [1. 1. 1.]]
[0 2 4 6 8]
[0.   0.25 0.5  0.75 1.  ]

Unlike a list, every element of a NumPy array has the same data type (here 64-bit integers, int64). That is what makes arrays compact and fast. np.zeros and np.ones create floating-point arrays by default, which is why they print with a decimal point. On some Windows setups, integer arrays may show int32 instead.

Shape, Size and Dimensions

results = np.array([
    [72, 81, 65],      # each row is a student
    [58, 64, 70],      # each column is a course
    [90, 77, 88],
    [45, 52, 60],
])
print(results.shape, results.ndim, results.size)

Output:

(4, 3) 2 12

The array has shape (4, 3): 4 rows and 3 columns, in 2 dimensions, with 12 elements in total.

Vectorised Arithmetic

Arithmetic operators work element by element on whole arrays, with no loop needed. This is called vectorisation.

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

print(a + b)
print(b / a)
print(a ** 2)
print(a * 3 + 1)
print(np.sqrt(b))

Output:

[11 22 33 44]
[10. 10. 10. 10.]
[ 1  4  9 16]
[ 4  7 10 13]
[3.16227766 4.47213595 5.47722558 6.32455532]

Compare this with plain Python, where [1, 2] + [3, 4] joins the lists into [1, 2, 3, 4] instead of adding them, and [1, 2] * 3 repeats the list. NumPy functions such as np.sqrt, np.sin, np.exp and np.log also work on every element at once.

How much faster?

import time

n = 1_000_000
xs = list(range(n))
arr = np.arange(n)

t0 = time.perf_counter()
squares_list = [x * x for x in xs]
t1 = time.perf_counter()
squares_arr = arr * arr
t2 = time.perf_counter()

print(squares_list[-1] == squares_arr[-1])
print("NumPy is faster:", (t2 - t1) < (t1 - t0))

Output:

True
NumPy is faster: True

On a typical laptop, the NumPy version is tens of times faster. The exact speed-up varies from machine to machine, so it is not printed here.

Indexing and Slicing

Indexing works like lists, but a 2D array takes a row and a column separated by a comma, and slices can be used in each dimension.

print(results[0])          # first row (student 0)
print(results[2, 1])       # row 2, column 1
print(results[:, 0])       # every row, column 0 (the first course)
print(results[1:3, 1:])    # rows 1-2, columns 1 onwards
print(results[-1, -1])     # last row, last column

Output:

[72 81 65]
77
[72 58 90 45]
[[64 70]
 [77 88]]
60

A slice of a NumPy array is a view of the same data, not a copy. Changing the slice changes the original array. Use .copy() when you need an independent copy.

Boolean Masks: Filtering

A comparison gives an array of True and False values, called a mask, which can then select elements.

print(scores >= 60)
print(scores[scores >= 60])               # keep only the passes
print(np.sum(scores >= 60), "passed")     # True counts as 1
print(scores[(scores >= 50) & (scores < 80)])

grades = np.where(scores >= 70, "A", "below A")
print(grades)

Output:

[ True False  True False  True]
[72 90 81]
3 passed
[72 58]
['A' 'below A' 'A' 'below A' 'A']

Combine conditions with & (and), | (or) and ~ (not), with each condition in brackets. Python's and and or do not work on arrays.

Statistics and Axes

print(scores.mean(), scores.max(), scores.min(), scores.sum())
print(round(scores.std(), 2), np.median(scores))

print(results.mean(axis=0))      # average of each column (each course)
print(results.mean(axis=1))      # average of each row (each student)
print(results.argmax(axis=0))    # which student topped each course

Output:

69.2 90 45 346
16.07 72.0
[66.25 68.5  70.75]
[72.66666667 64.         85.         52.33333333]
[2 0 2]

The axis argument says which dimension to collapse. axis=0 works down the rows, giving one result per column, and axis=1 works across the columns, giving one result per row. argmax returns the position of the largest value instead of the value itself.

std() calculates the population standard deviation by default. For a sample standard deviation, which divides by n − 1, use std(ddof=1).

Broadcasting

When two arrays have different shapes, NumPy can often stretch the smaller one to match. This is called broadcasting. A single number is broadcast to every element, and a row can be broadcast down every row of a 2D array.

units = np.array([3, 2, 1])            # credit units of the three courses
weighted = results * units             # each column multiplied by its units
print(weighted[0])

curve = results + np.array([5, 0, 2])  # add a different bonus to each course
print(curve[3])

normalised = (results - results.mean(axis=0)) / results.std(axis=0)
print(normalised.round(2)[0])

Output:

[216 162  65]
[50 52 62]
[ 0.34  1.1  -0.54]

The last line standardises each course's marks into z-scores: how many standard deviations each mark is above or below that course's mean.

Reshaping and Combining

nums = np.arange(1, 13)
grid = nums.reshape(3, 4)
print(grid)
print(grid.T.shape)                    # .T is the transpose
print(grid.flatten()[:5])
print(np.vstack([a, b]))               # stack arrays as rows

Output:

[[ 1  2  3  4]
 [ 5  6  7  8]
 [ 9 10 11 12]]
(4, 3)
[1 2 3 4 5]
[[ 1  2  3  4]
 [10 20 30 40]]

The new shape in reshape must contain the same number of elements: 12 = 3 × 4. Passing -1 for one dimension lets NumPy work it out, so reshape(-1, 2) gives 6 rows of 2.

Linear Algebra

NumPy handles matrices directly. The @ operator performs matrix multiplication (whereas * multiplies element by element), and np.linalg provides determinants, inverses and equation solvers.

A = np.array([[2, 1],
              [1, 3]])
B = np.array([[1, 0],
              [2, 1]])
print(A @ B)
print(round(np.linalg.det(A), 6))
print(np.linalg.inv(A).round(2))

# Solve the simultaneous equations 2x + y = 5 and x + 3y = 10
coeffs = np.array([[2, 1], [1, 3]])
consts = np.array([5, 10])
print(np.linalg.solve(coeffs, consts))

Output:

[[4 1]
 [7 3]]
5.0
[[ 0.6 -0.2]
 [-0.2  0.4]]
[1. 3.]

The solution is x = 1 and y = 3. Check: 2(1) + 3 = 5, and 1 + 3(3) = 10. np.linalg.solve solves systems with hundreds or thousands of unknowns just as easily, which is how engineers analyse circuits and structures.

Random Numbers and Simulation

rng = np.random.default_rng(seed=42)     # a seed makes the results repeatable

rolls = rng.integers(1, 7, size=60000)   # 60,000 dice rolls
counts = np.bincount(rolls)[1:]          # how many of each face, 1 to 6
print(len(counts), counts.sum(), bool(np.all(abs(counts - 10000) < 500)))

heights = rng.normal(loc=165, scale=8, size=100_000)
print(round(heights.mean()), round(heights.std()))
within = np.mean(abs(heights - 165) < 8)
print(round(within, 2))

Output:

6 60000 True
165 8
0.68

Each face of the die comes up close to 10,000 times, and about 68% of the simulated heights lie within one standard deviation of the mean, just as the normal distribution predicts. Simulations like this let you check probability results by experiment.

Common Mistakes with NumPy

  • Using * for matrix multiplication. It multiplies element by element. Use @.
  • Using and or or with arrays. Use & and |, with brackets around each condition.
  • Forgetting that slices are views, so that changing a slice changes the original array.
  • Looping over elements in Python instead of using vectorised operations, which throws away NumPy's speed.
  • Mixing up axis=0 and axis=1. axis=0 gives one result per column.

NumPy Basics: Practice Exercises

  1. Create an array of the temperatures 30, 32, 35, 31 and 29 (in °C), and convert them all to Fahrenheit.
  2. Using the results array, find the highest mark of each student.
  3. Count how many marks in results are below 60.
  4. Solve the system x + y + z = 6, 2x − y + z = 3 and x + 2y − z = 2 with np.linalg.solve.
  5. Create the 5 × 5 multiplication table using broadcasting.

Sample solutions:

temps = np.array([30, 32, 35, 31, 29])
print(temps * 9 / 5 + 32)

print(results.max(axis=1))

print(int(np.sum(results < 60)))

M = np.array([[1, 1, 1], [2, -1, 1], [1, 2, -1]])
print(np.linalg.solve(M, np.array([6, 3, 2])).round(6))

n = np.arange(1, 6)
print(n.reshape(-1, 1) * n)

Output:

[86.  89.6 95.  87.8 84.2]
[81 70 90 60]
3
[1. 2. 3.]
[[ 1  2  3  4  5]
 [ 2  4  6  8 10]
 [ 3  6  9 12 15]
 [ 4  8 12 16 20]
 [ 5 10 15 20 25]]

In the last solution, a column of shape (5, 1) times a row of shape (5,) broadcasts to a full 5 × 5 grid.

Conclusion

NumPy stores numbers in fast, single-type arrays and applies operations to whole arrays at once. Create arrays with np.array, arange, linspace, zeros and ones, and inspect them with shape and dtype. Use vectorised arithmetic instead of loops, index 2D arrays with [row, column], and filter with Boolean masks. Summarise with mean, std and argmax along an axis, and rely on broadcasting to combine arrays of different shapes. Multiply matrices with @ and solve equations with np.linalg.solve. NumPy is the foundation for pandas, plotting and machine learning in Python.

Related Tutorials