This is one of the 100 recipes of the IPython Cookbook, the definitive guide to high-performance scientific computing and data science in Python.

4.8. Making efficient selections in arrays with NumPy

In [ ]:
import numpy as np
In [ ]:
id = lambda x: x.__array_interface__['data'][0]

We create a large array.

In [ ]:
n, d = 100000, 100
In [ ]:
a = np.random.random_sample((n, d)); aid = id(a)

Array views and fancy indexing

We take a selection using two different methods: with a view and with fancy indexing.

In [ ]:
b1 = a[::10]
b2 = a[np.arange(0, n, 10)]
In [ ]:
np.array_equal(b1, b2)

The view refers to the original data buffer, whereas fancy indexing yields a copy.

In [ ]:
id(b1) == aid, id(b2) == aid

Fancy indexing is several orders of magnitude slower as it involves copying a large array. Fancy indexing is more general as it allows to select any portion of an array (using any list of indices), not just a strided selection.

In [ ]:
%timeit a[::10]
In [ ]:
%timeit a[np.arange(0, n, 10)]

Alternatives to fancy indexing: list of indices

Given a list of indices, there are two ways of selecting the corresponding sub-array: fancy indexing, or the np.take function.

In [ ]:
i = np.arange(0, n, 10)
In [ ]:
b1 = a[i]
b2 = np.take(a, i, axis=0)
In [ ]:
np.array_equal(b1, b2)
In [ ]:
%timeit a[i]
In [ ]:
%timeit np.take(a, i, axis=0)

Using np.take instead of fancy indexing is faster.

Note: Performance of fancy indexing has been improved in recent versions of NumPy; this trick is especially useful on older versions of NumPy.

Alternatives to fancy indexing: mask of booleans

Let's create a mask of booleans, where each value indicates whether the corresponding row needs to be selected in x.

In [ ]:
i = np.random.random_sample(n) < .5

The selection can be made using fancy indexing or the np.compress function.

In [ ]:
b1 = a[i]
b2 = np.compress(i, a, axis=0)
In [ ]:
np.array_equal(b1, b2)
In [ ]:
%timeit a[i]
In [ ]:
%timeit np.compress(i, a, axis=0)

Once again, the alternative method to fancy indexing is faster.

You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).

IPython Cookbook, by Cyrille Rossant, Packt Publishing, 2014 (500 pages).