CODEWITHSUNDEEP
pythonarraysnumpyindexing
Alejandro Sazo
Alejandro Sazo

Reputation: 795

Multiple lists of indices for one array in numpy

In normal situations a list with integers can be used as indices for an array. Let's say

arr = np.arange(10)*2
l = [1,2,5]
arr[l] # this gives np.array([2,4,10])

Instead of one list of indices, I have several, with different lenghts, an I want to get arr[l] for each sublist in my list of indices. How can I achieve this without an sequential approach (using a for), or better, using less time than using a for using numpy?

For example:

lists = [[1,2,5], [5,6], [2,8,4]]
arr = np.arange(10)*2
result = np.array([[2,4,10], [10, 12], [4,16,8]]) #this is after the procedure I want to get

Upvotes: 1

Views: 1381

Answers (1)

eickenberg
eickenberg

Reputation: 14377

It depends on the size of your lists whether this makes sense. One option is to concatenate them all, do the slicing and then redistribute into lists.

lists = [[1,2,5], [5,6], [2,8,4]]
arr = np.arange(10)*2

extracted = arr[np.concatenate(lists)]

indices = [0] + list(np.cumsum(map(len, lists)))
result = [extracted[indices[i]:indices[i + 1]] for i in range(len(lists))]

Or, taking into account @unutbu's comment:

result = np.split(extracted, indices[1:-1])

Upvotes: 2

Related Questions