CODEWITHSUNDEEP
pythonnumpy
Fernando ardana Llamas
Fernando ardana Llamas

Reputation: 11

How to select elements on given axis base on the value of another array

I am triying to solve the following problem in a more numpy-friendly way (without loops):

We want the a NxM matrix (R) with R[i,j] = D[k,i,j] being k=g[i,j] A loop base solution is:

def getVals(g, d):
    arr=np.zeros(g.shape)
    for row in range(g.shape[0]):
        for column in range(g.shape[1]):
            arr[row,column]=d[g[row,column],row,column]
    return arr

Upvotes: 1

Views: 36

Answers (3)

obchardon
obchardon

Reputation: 10792

You could also use np.take_along_axis

Then you can simply extract your values along one specific axis:

# Example input data:
G = np.random.randint(0,3,(4,5))    #   4x5 array
D = np.random.randint(0,9,(3,4,5))  # 3x4x5 array

# Get the results:
R = np.take_along_axis(D,G[None,:],axis=0)

Since G should have the same number of dimension as D, we simply add a new dimension to G with G[None,:].

Upvotes: 0

sprico27
sprico27

Reputation: 42

Here's my try (I assume g and d are Numpy Ndarrays):

def getVals(g, d):
    m,n = g.shape
    indexes = g.flatten()*m*n + np.arange(m*n)
    arr = d.flatten()[indexes].reshape(m,n)
    return arr

So if

d = [[[96, 89, 51, 40, 51],
    [31, 72, 39, 77, 33]],

   [[34, 11, 54, 86, 73],
    [12, 21, 74, 39, 14]],

   [[14, 91, 38, 77, 97],
    [44, 55, 93, 88, 55]]]

and

g = [[2, 1, 2, 1, 1],
   [0, 2, 0, 0, 2]]

then you are going to get

arr = [[14, 11, 38, 86, 73],
   [31, 55, 39, 77, 55]]

Upvotes: 0

Quang Hoang
Quang Hoang

Reputation: 150735

Try with ogrid and advanced indexing:

x,y = np.ogrid[:N,:M]
out = D[G, x[None], y[None]]

Test:

N,M=4,5
G = np.random.randint(0,3, (N,M))
D = np.random.rand(3,N,M)

np.allclose(getVals(G,D), D[G, x[None], y[None]])
# True

Upvotes: 1

Related Questions