Indexing: A few handy ways to access NumPy arrays

9 minute read

The following code snippets should serve as an (incomplete) cheat sheet for accessing NumPy arrays. All examples expect an import numpy as np.

Basic access

NumPy arrays can be accessed just like lists with array[start:stop:step]

a = np.array([1,2,3,4], int)
# => array([1, 2, 3, 4])
a[2]
# => 3
a[:2]
# => array([1, 2])
a[::2]
# => array([1, 3])

When working with multidimensional arrays, a comma can be used to access values for the different axes:

b = np.array([[1,2,3,4],[5,6,7,8]], int)
# => array([[1, 2, 3, 4],
#          [5, 6, 7, 8]])
b[0,2]
# => 3
b[0,:]
# => array([1, 2, 3, 4])
b[0,::2]
# => array([1, 3])

Negative values can be used to access the end of the array

b[0,-1]
# => 4

Access using integer index arrays

To get a subset of an array via the indices, integer arrays can be used. Say we want to access the first, third and fifth element of a one-dimensional array:

a = np.array([1,2,3,4,5,6], int)
# => array([1, 2, 3, 4, 5, 6])
selector = np.array([0,2,4],int)
a[selector]
# => array([1, 3, 5])

This also works for multi-dimensional arrays when one-dimensional arrays are passed for each axis.

a = np.array([[1,2,3],[4,5,6]], int)
# => array([[1, 2, 3],
#          [4, 5, 6]])
sel1 = np.array([0,1,0], int)
sel2 = np.array([1,2,1], int)
a[sel1,sel2]
# => array([2, 6, 2])

In the previous example the two selector (index) arrays are used to access…

  • index 0 for axis 0, index 1 for axis 1 (==2)
  • index 1 for axis 0, index 2 for axis 1 (==6)
  • index 0 for axis 0, index 1 for axis 1 (==2)

To better mentally visualize this, the 2D array axis 0 could be thought of as the “rows”, axis 1 the “columns”.

Selecting across dimensions is also possible. By using the take function, an axis can be specified.

a = np.array([[1,2,3],[4,5,6]], int)
# => array([[1, 2, 3],
#          [4, 5, 6]])
sel = np.array([0,1], int)
a.take(sel, axis=0)
# => array([[1, 2, 3],
#          [4, 5, 6]])
a.take(sel, axis=1)
# => array([[1, 2],
#          [4, 5]])

Access using boolean arrays

To select using a boolean array, we can do the following:

a = np.array([[1,2,3,4],[5,6,7,8]], int)
# => array([[1, 2, 3, 4],
#          [5, 6, 7, 8]])
b = a>3
# => array([[False, False, False,  True],
#          [ True,  True,  True,  True]], dtype=bool)
a[b]
# => array([4, 5, 6, 7, 8])

The same works for boolean expressions like a logical_or, logical_and, logical_not, logical_xor:

a[np.logical_or(a<=4,a>=7)]
# => array([1, 2, 3, 4, 7, 8])

Or, as another example, for getting elements which are not (~) NaN:

a = np.array([[1,2], [3,np.nan]], float)
# => array([[  1.,   2.],
#          [  3.,  nan]])
a[~np.isnan(a)]
# => array([ 1.,  2.,  3.])

More information

There are, of course, many more ways to juggle around with NumPy arrays. A more complete introduction on indexing can be found at: https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html

Tags:

Categories:

Updated:

Like to comment? Feel free to send me an email or reach out on Twitter.

Did this or another article help you? If you like and can afford it, you can buy me a coffee (3 EUR) ☕️ to support me in writing more posts. In case you would like to contribute more or I helped you directly via email or coding/troubleshooting session, you can opt to give a higher amount through the following links or adjust the quantity: 50 EUR, 100 EUR, 500 EUR. All links redirect to Stripe.