CODEWITHSUNDEEP
python
Stumpp
Stumpp

Reputation: 289

Numpy retrieve a range of elements from a given set of indexes

Lets say we have a numpy array = [10,11,12,13,14,15,16,17,18,19]. If you want to change the elements with index = [1,3,5] in this array you can simply call array[index] which will give you an array of the elements [11,13,15].

Now if we still have array = [10,11,12,13,14,15,16,17,18,19] and I want to pick out a sequence of, ex. length 2, of numbers in this array based on a array of index = [1,3,5], is there a nice way in numpy to do this by doing something like array[index:(index+2)] returning [[11,12],[13,14],[15,16]]? I could obviously do this easily with for-loops but I wonder if there is a faster way?

Upvotes: 0

Views: 41

Answers (2)

Aryerez
Aryerez

Reputation: 3495

I think this is the best way:

array = [10,11,12,13,14,15,16,17,18,19]
index = [1, 3, 5]
count = 2
[array[num:num + count] for num in index]

Upvotes: 0

John Zwinck
John Zwinck

Reputation: 249103

I'd do it this way:

array[index.reshape(-1,1) + np.arange(2)]

The argument to arange() determines how many elements to take starting from each index, and it's OK if they overlap. For example with arange(4) the result is:

  [[11, 12, 13, 14],
   [13, 14, 15, 16],
   [15, 16, 17, 18]]

Upvotes: 1

Related Questions