AlexT
AlexT

Reputation: 609

Python - How to extract elements from an array based on an array of indices?

Let's say I have a list of elements X and one of indices Y.

X = [1, 2, 3, 4, 5, 6, 7]
Y = [0, 3, 4]

Is there a function in Python that allows one to extract elements from X based on the indices provided in Y? After execution, X would be:

X = [1, 4, 5]

Upvotes: 1

Views: 12505

Answers (3)

jpp
jpp

Reputation: 164623

You can use list.__getitem__ with map:

X = [1, 2, 3, 4, 5, 6, 7]
Y = [0, 3, 4]

res = list(map(X.__getitem__, Y))  # [1, 4, 5]

Or, if you are happy to use a 3rd party library, you can use NumPy:

import numpy as np

X = np.array([1, 2, 3, 4, 5, 6, 7])
res = X[Y]  # array([1, 4, 5])

Upvotes: 1

sacuL
sacuL

Reputation: 51335

The list comprehension provided by @Prune is the way to go in pure python. If you don't mind numpy, it might be easier just use their indexing scheme:

import numpy as np
>>> np.array(X)[Y]
array([1, 4, 5])

Upvotes: 2

Prune
Prune

Reputation: 77837

X = [X[index] for index in Y]

This is a list comprehension; you can look up that topic to learn more.

Upvotes: 8

Related Questions