CODEWITHSUNDEEP
numpy
YohanRoth
YohanRoth

Reputation: 3253

How to return list of items from an array given list on indices in numpy

I have an array arr and a list of indices that I want to get indices. I want to get subset of array corresponding to items in indices and the complement of that.

For example

for 

arr = np.asarray([2, 4, 1, 5, 6])
indices = np.asarray([2, 4])

I would get

[1, 6] and [2, 4, 5]

Thanks

Upvotes: 0

Views: 42

Answers (1)

Tarifazo
Tarifazo

Reputation: 4343

Using np.isin or np.in1d (using masks):

arr = np.asarray([2, 4, 1, 5, 6])
indices = np.asarray([2, 4])
m = np.in1d(np.arange(len(arr)), indices)
arr1, arr2 = arr[m], arr[~m]
arr1, arr2
>>array([1, 6]), array([2, 4, 5])

Alternatively, using np.setdiff1d for the complementary part (can be faster for larger arrays and indices):

arr1 = arr[indices]
arr2 = arr[np.setdiff1d(np.arange(len(arr)), indices, True)]

Upvotes: 2

Related Questions