Multiple lists of indices for one array in numpy

Question

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

I don't think you can avoid using an iterative approach, since your result is a list, not a Numpy type (and you can't make it be a numpy type since the sublist lengths aren't equal). — BrenBarn
– BrenBarn, Commented Aug 24, 2014 at 21:00
the result can be a numpy array, I'm going to edit to show this — Alejandro Sazo
– Alejandro Sazo, Commented Aug 24, 2014 at 21:00
Although result is a NumPy array, it has dtype object, which is the worst possible type of NumPy array. It is an array of pointers to Python objects. Using such an array is typically slower than using a Python list. — unutbu
– unutbu, Commented Aug 24, 2014 at 21:04
@AlejandroSazo what about something like this: [arr[lists[k]] for k in range(len(lists))] — Dalek
– Dalek, Commented Aug 24, 2014 at 21:06
@Dalek, That is what I am currently using, I wanted something like vectorized, or faster — Alejandro Sazo
– Alejandro Sazo, Commented Aug 24, 2014 at 21:07

eickenberg · Accepted Answer · 2014-08-24 21:15:04Z

2

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

edited Aug 24, 2014 at 21:15

answered Aug 24, 2014 at 21:00

eickenberg

14.4k1 gold badge52 silver badges52 bronze badges

Sign up to request clarification or add additional context in comments.

8 Comments

unutbu Over a year ago

Also possible is: indices = np.cumsum(map(len, lists)), np.split(extracted, indices[:-1]).

eickenberg Over a year ago

That is amazingly concise - I didn't know np.split.

Alejandro Sazo Over a year ago

I am trying to test the answer but I didn't know I have some empty lists, so the concatenation is made of floats instead of the desired ints:(

eickenberg Over a year ago

I tried adding an empty list to lists and for me the code runs through and returns an empty array at that point, so it works at least on my config. Where do the floats come from? I don't see how they could be introduced.

Alejandro Sazo Over a year ago

with the same example lists and adding an empty list, the concatenation throws an error:

>>> lists = [[], [1,2,5], [5,6], [2,8,4]] >>> arr[np.concatenate(lists)] Traceback (most recent call last):   File "<stdin>", line 1, in <module> IndexError: arrays used as indices must be of integer (or boolean) type

|

Collectives™ on Stack Overflow

Multiple lists of indices for one array in numpy

1 Answer 1

8 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

8 Comments

Your Answer

Sign up or log in

Post as a guest

Related