CODEWITHSUNDEEP
pythonarraysnumpy
Cupitor
Cupitor

Reputation: 11637

Vectorised code for selecting elements of a 2D array

Well the following code obviously returns the element in position ind in matrix:

def select_coord(a,ind):
    return a[ind]

However I don't know how to vectorise this. In other words:

b=np.asarray([[2,3,4,5],[7,6,8,10]])
indices=np.asarray([2,3])
select_coord(b,indices)

Should return [4,10].

Which can be written with a for loop:

def new_select_record(a,indices):
    ret=[]
    for i in range a.shape[0]:
        ret.append(a[indices[i]])
    return np.asarray(ret)

Is there a way to write this in a vectorised manner?

Upvotes: 1

Views: 44

Answers (2)

falsetru
falsetru

Reputation: 369164

To get b[0, 2], b[1, 3]:

>>> import numpy as np
>>> b = np.array([[2,3,4,5], [7,6,8,10]])
>>> indices = np.array([2, 3])
>>> b[np.arange(len(indices)), indices]
array([ 4, 10])

Upvotes: 1

CT Zhu
CT Zhu

Reputation: 54340

how about: np.diag(b[:,[2,3]])?

Upvotes: 1

Related Questions