NumPy indexing explained
By: Àlex Escolà Nixon
Tags: Numpy, Data Science, Artificial Intelligence, Machine Learning

NumPy indexing explained

Photo by Jean-Louis Paulin on Unsplash

NumPy is the universal standard for working with Numerical data in Python. Multidimensional NumPy arrays are extensively used in Pandas, SciPy, Scikit-Learn, scikit-image, which are some of the main data science and scientific Python packages. So a good understanding of NumPy is crucial if we are working with these tools! And so it is to understand its indexing methods, which we will go through in this post.

Note that a key aspect or motivation in using NumPy, is that it enables to manipulate large quantities of numerical data in a vectorised way, that is, avoiding the use of inefficient python loops. The huge advantage of vectorising your code is that you’re pushing all looping down to C level, which is much faster.

Now let’s get to the topic of the post, which is indexing NumPy arrays. Indexing with NumPy can be a little tough and counter-intuitive without a good understanding of its basic rules. Here we’ll dive into the different indexing methods that can be used, which are:

- Basic Slicing and Indexing - Advanced Indexing (including both integer and boolean array indexing)

Find a notebook version of this post along with all explanations here.

Basic slicing

NumPy’s basic slicing is an extension of Python’s basic slicing concept extended to N dimensions.

It essentially allows you to take slices of an array along its dimensions using basic slicing notation, i.e start:stop:step. If you are not too familiar with Python’s basic slicing notation you can check this post in stack overflow which makes it very clear. Let’s see with an example:

array([['A', 'B', 'C', 'D', 'E'],

['F', 'G', 'H', 'I', 'J'],

['K', 'L', 'M', 'N', 'O'],

['P', 'Q', 'R', 'S', 'T'],

['U', 'V', 'W', 'X', 'Y']])

Yes, I know z is missing… Nothing against it, but we need a homogeneous matrix. This example array has a shape of (5, 5), so two dimensions. We can take slices along its dimensions (I’ll be using the equivalent term axes) by following the mentioned notation, for instance:

a[:2, 1:4] array([['B', 'C', 'D'],

['G', 'H', 'I']])

So far so good. We’ve basically sliced the array along the first axis to get up to the second row starting at index 0 (note that the stop index is not included!), and along the second to get columns one to three. Think of the axes just as the (x,y,z) dimensions of a matrix. Where a third dimension can be simply though of as stacking multiple 2D arrays together:

Okay, so now we know we can slice an n-dimensional array using python’s slice notation. But what if you had to take values from a given axis using a list of indexes?

Well, here’s where advanced indexing comes into play ↓

Advanced indexing

Let’s just go with an example of what an intuitive way of doing this would be:

a_indexed = a[[4,3,1], [2,4,0]]

Okay, so following the same logic as in basic slicing, one could expect to get all values that fall into the above indexes along the different axes of the ndarray. Namely all values in rows 4, 3 and 1 which are in columns 2, 4 and 0 respectively, so the intersecting values that are highlighted bellow:

Is that so? Let’s check:

print(a_indexed)

array(['W', 'T', 'F'])

Clearly not… (example inspired by this very instructive talk by Jaime Fernández :) ) What happened there? Why are we getting a one-dimensional result?

This is because advanced indexing follows a different set of rules. A good way to think about it, is that when using basic slicing, we are indexing on a grid which is defined by the slices we take on each dimension. Whereas using advanced indexing can be thought of as specifying a set of (x,y) coordinates of the values we want to retrieve.

In the above case indexing with [4,3,1],[2,4,0] , will behave as indexing on (4,2) , (3,4) and (1,0) respectively. So each index we specify in each dimension, will be combined with the corresponding indexes from the other dimensions.

More generally, when using advanced indexing, we must take into account these two main aspects (from the docs):

The indexing arrays represent a number of indexes into that dimension

Which basically means that we will retrieve as many elements from a given axis as indexes specified in the indexing array into that dimension. In the above case, we’ve retrieved 3 elements specified by the length 3 indexing arrays.

The result shape is identical to the (broadcast) indexing array shapes

To understand the above we need to dive a little into broadcasting…

Going deep into broadcasting could easily lead to an another entire blog post, so i’ll just try to cover the very essentials… From the docs:

broadcasting is defined as a term describing how NumPy treats arrays with different shapes during arithmetic operations. Subject to certain constraints, the smaller array is “broadcast” across the larger array so that they have compatible shapes

So basically, when some operation involving arrays with different shapes is performed, NumPy tries to make their shapes compatible before the operation takes place. Let’s take a look at some visual examples:

In the first case for instance, a one-dimensional NumPy array is added to an integer. So under the hood, before the addition takes place NumPy will broadcast the smaller shaped array across the larger, i.e it will replicate its value until its shape becomes compatible with that of the larger array.

And how do we know if two shapes are compatible?

We can find the general broadcasting rules defined in the docs. It is stated that two dimensions are compatible when