CODEWITHSUNDEEP
pythonnumpy
hl037_
hl037_

Reputation: 3867

Fastest way to extract elements from a multi-dimensionnal array at specified array of indices

Suppose I have an n-D array A, with an arbitrary n:

a = np.arange(3*4*5*...).reshape((3, 4, 5, ...))

...And I have an array of indices (exactly this shape or transposed):

ix = np.array([
  [0,1,1,...],
  [0,0,1,...],
  [1,0,1,...],
])

What is the fastest way of getting an array of the elements from a at indices from ix ?

(e.g. the fastest way of obtaining :

def take_at(a, ix):
  res = np.empty((ix.shape[0],*a.shape[ix.shape[1]:]))
  for i in range(ix.shape[0]):
    res[i,...] = a[(*ix[i],)]
  return res

...Using (if possible) only numpy vectorized function / loop / operations ?)

simple copy-pastable example to test : :

a = np.arange(3*4*5).reshape((3,4,5))

ix = np.array([
  [0,1,1],
  [0,0,1],
  [1,0,1],
])

def take_at(a, ix):
  res = np.empty((ix.shape[0],*a.shape[ix.shape[1]:]), dtype=a.dtype)
  for i in range(ix.shape[0]):
    res[i,...] = a[(*ix[i],)]
  return res

take_at(a, ix)

Upvotes: 0

Views: 36

Answers (2)

Daniel F
Daniel F

Reputation: 14399

You need to pack ix into a tuple:

take_at(a, ix)
Out[]: array([ 6,  1, 21])

a[tuple(ix.T)]
Out[]: array([ 6,  1, 21])

Upvotes: 1

Quang Hoang
Quang Hoang

Reputation: 150735

You can do:

a[ix[:,0],ix[:,1],ix[:,2]]

Output:

array([ 6,  1, 21])

Upvotes: 0

Related Questions