CODEWITHSUNDEEP
pythonnumpyoptimization
Roman
Roman

Reputation: 131058

How to extract a list of elements given by their indices from a numpy array efficiently?

I have a multidimensional numpy array and I would like to take some of its elements to construct a one dimensional array. The elements that I need to take are given by their indices, for example:

inds = [(0,0), (0,1), (1,1), (1,0), (0,2)]

I solve it in a straightforward way:

ls = [] 
for i, j in inds:
   ls += [a[i,j]]

It gives the desired result. However, I have realized that this solution is too slow for my purposes. Are there a possibility to do the same in a more efficient way?

Upvotes: 3

Views: 309

Answers (1)

Warren Weckesser
Warren Weckesser

Reputation: 114811

numpy arrays can be indexed with sequences (and, more generally, numpy arrays).

For example, here's my array a

In [19]: a
Out[19]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14]])

i and j hold the sequences of the first and second coordinates of your inds array:

In [20]: i
Out[20]: [0, 0, 1, 1, 0]

In [21]: j
Out[21]: [0, 1, 1, 0, 2]

You can use these to pull the corresponding values out of a:

In [22]: a[i, j]
Out[22]: array([0, 1, 6, 5, 2])

If you already have inds in your code, you can separate the list of tuples into i and j using zip:

In [23]: inds
Out[23]: [(0, 0), (0, 1), (1, 1), (1, 0), (0, 2)]

In [24]: i, j = zip(*inds)

In [25]: i
Out[25]: (0, 0, 1, 1, 0)

In [26]: j
Out[26]: (0, 1, 1, 0, 2)

Or, if inds is an array with shape (n, 2), like so:

In [27]: inds = np.array(inds)

In [28]: inds
Out[28]: 
array([[0, 0],
       [0, 1],
       [1, 1],
       [1, 0],
       [0, 2]])

you can simply assign the transpose of inds to i, j:

In [33]: i, j = inds.T

In [34]: i
Out[34]: array([0, 0, 1, 1, 0])

In [35]: j
Out[35]: array([0, 1, 1, 0, 2])

In [36]: a[i, j]
Out[36]: array([0, 1, 6, 5, 2])

Upvotes: 6

Related Questions