Indexing: A few handy ways to access NumPy arrays

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