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

Question

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?

Check out the [documentation on boolean indexing](https://numpy.org/devdocs/reference/arrays.indexing.html#boolean-array-indexing). — Bertil Johannes Ipsen, Jun 02 '20 at 15:25
Does this answer your question? [Iterate over numpy with index (numpy equivalent of python enumerate)](https://stackoverflow.com/questions/42082607/iterate-over-numpy-with-index-numpy-equivalent-of-python-enumerate) — DaVinci, Jun 02 '20 at 15:29
Actually, @Quang Hoang's answer solves my problem efficiently ^^. I work with very large matrix, so I think the solution in the linked answer is not as efficient. — Akira, Jun 02 '20 at 15:31

score 2 · Accepted Answer · answered Jun 02 '20 at 15:25

2

I think you want boolean indexing:

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

answered Jun 02 '20 at 15:25

Quang Hoang

score 0 · Answer 2 · answered Jun 02 '20 at 15:31

0

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]]

answered Jun 02 '20 at 15:31

match

2 Answers2