Indexing numpy array with another numpy array

Question

Suppose I have

a = array([[1, 2],
           [3, 4]])

and

b = array([1,1])

I'd like to use b in index a, that is to do a[b] and get 4 instead of [[3, 4], [3, 4]]

I can probably do

a[tuple(b)]

Is there a better way of doing it?

Thanks

Answer 1

According the NumPy tutorial, the correct way to do it is:

a[tuple(b)]

Answer 2

Suppose you want to access a subvector of a with n index pairs stored in blike so:

b = array([[0, 0],
       ...
       [1, 1]])

This can be done as follows:

a[b[:,0], b[:,1]]

For a single pair index vector this changes to a[b[0],b[1]], but I guess the tuple approach is easier to read and hence preferable.

Answer 3

The above is correct. However, if you see an error like:

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

You may have your index array in floating type. Change it to something like this:

arr[tuple(a.astype(int))]

Answer 4

For performance, use

a[b[0], b[1]]
a[b[:,0], b[:,1]]             # when b is a 2d array

unless d = tuple(b.T) is readily available (a[d] is the fastest)

---

time bench shows a[b[:,0], b[:,1]] is significantly faster than a[tuple(c)] both when b is 1d and 2d (omitted), where c == b.T

import numpy as np
import time

a = np.array([[1, 2], [3, 4]])s
b = np.array([1,1])
c = b.T                       # a[b[:,0], b[:,1]] == a[tuple(c)] when b is 2d
assert tuple(c) == tuple(b)
d = tuple(b)

t0 = time.time()
for _ in range(10000000):
    a[tuple(c)]
t1 = time.time()
print(t1 - t0)                # 7.781806468963623

t0 = time.time()
for _ in range(10000000):
    a[b[0], b[1]]
t1 = time.time()
print(t1 - t0)                # 2.8317997455596924

t0 = time.time()
for _ in range(10000000):
    a[*b]
t1 = time.time()
print(t1 - t0)                # 7.80025315284729

t0 = time.time()
for _ in range(10000000):
    a[d]
t1 = time.time()
print(t1 - t0)                # 1.3347711563110352