CODEWITHSUNDEEP
python-3.xnumpysubset
Akira
Akira

Reputation: 2870

How to extract elements in an array with regard to an index?

I have a row A = [0 1 2 3 4] and an index I = [0 0 1 0 1]. I would like to extract the elements in A indexed by I, i.e. [2, 4].

My attempt:

import numpy as np
A = np.array([0, 1, 2, 3, 4])
index = np.array([0, 0, 1, 0, 1])
print(A[index])

The result is not as I expected:

[0 0 1 0 1]

Could you please elaborate on how to achieve my goal?

Upvotes: 0

Views: 523

Answers (2)

match
match

Reputation: 11060

A non-numpy way to achieve this, in case its useful - it uses zip to combine each pair of elements, and returns the first if the second is true:

[x[0] for x in zip(a, i) if x[1]]

Upvotes: 0

Quang Hoang
Quang Hoang

Reputation: 150735

I think you want boolean indexing:

A[index.astype(bool)]
# array([2, 4])

Upvotes: 2

Related Questions