NumPy: Arrays and Vectorized Computation

NumPy: arrays and vectorized computation

NumPy is the numerical computing library on which pandas, scikit-learn and most data tools in Python are built. Its basic building block is the array (ndarray): a grid of values of a single type, stored compactly in memory.

Creating an array

import numpy as np   # universal alias

prix = np.array([49.0, 19.0, 189.0, 25.5])
print(prix.shape)    # (4,)     : 4 elements in one dimension
print(prix.dtype)    # float64  : all elements have the same type

np.arange(0, 10, 2)  # array([0, 2, 4, 6, 8])
np.zeros(3)          # array([0., 0., 0.])
np.linspace(0, 1, 5) # 5 evenly spaced values between 0 and 1

Computing without loops: vectorization

With a list, applying VAT requires a loop or a comprehension. With NumPy, you write the operation directly on the array:

prix_ttc = prix * 1.2        # each element is multiplied
quantites = np.array([3, 10, 1, 4])
ca = prix * quantites        # element-by-element multiplication
print(ca.sum())              # 628.0

This vectorized calculation runs in compiled code: on a million values, it is dozens of times faster than a Python loop. Combining an array and a plain number (prix * 1.2) is called broadcasting: NumPy “stretches” the number to the size of the array.

Be careful: with a list, [1, 2] * 2 gives [1, 2, 1, 2] (repetition), whereas with NumPy np.array([1, 2]) * 2 gives [2, 4].

Filtering with a boolean mask

A comparison on an array returns an array of booleans, called a mask:

masque = prix > 30
print(masque)              # [ True False  True False]
print(prix[masque])        # [ 49. 189.]  : only the Trues are kept
print(prix[(prix > 20) & (prix < 100)])   # [49.  25.5]
print(masque.sum())        # 2 : True counts as 1

To combine conditions, use & (and), | (or), ~ (not) and put each condition in parentheses. The and / or keywords do not work on arrays.

Two-dimensional arrays and axis

# ventes: 3 stores (rows) x 4 quarters (columns)
ventes = np.array([[10, 12, 9, 14],
                   [ 8,  7, 11, 10],
                   [15, 14, 13, 18]])
print(ventes.shape)          # (3, 4)
print(ventes[0, 2])          # 9 : row 0, column 2
print(ventes.sum(axis=1))    # [45 36 60] : total per store (the columns are collapsed)
print(ventes.mean(axis=0))   # average per quarter (the rows are collapsed)
print(ventes.max())          # 18 : over the whole array

Remember: axis=0 aggregates along the rows (one result per column), axis=1 aggregates along the columns (one result per row).

Key takeaways

• A NumPy array has a single type and a shape (shape).
• Operations are vectorized: no loop.
• Filter with masks and &, |, ~ in parentheses.