Ritvik Vij
Ritvik Vij

Reputation: 33

Slicing Numpy Array by 2 index arrays

If I have a set of indices stored in two Numpy arrays, my goal is to slice a given input array based on corresponding indices in those index arrays. For eg.

index_arr1 = np.asarray([2,3,4])
index_arr2 = np.asarray([5,5,6])

input_arr = np.asarray([1,2,3,4,4,5,7,2])

The output to my code should be [[3,4,4],[4,4],[4,5]] which is basically [input_arr[2:5], input_arr[3:5], input_arr[4:6]]

Can anybody suggest a way to solve this problem using numpy functions and avoiding any for loops to be as efficient as possible.

Upvotes: 0

Views: 92

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150805

Do you mean:

[input_arr[x:y] for x,y in zip(index_arr1, index_arr2)]

Output:

[array([3, 4, 4]), array([4, 4]), array([4, 5])]

Or if you really want list of lists:

[[input_arr[x:y].tolist() for x,y in zip(index_arr1, index_arr2)]

Output:

[[3, 4, 4], [4, 4], [4, 5]]

Upvotes: 2

Related Questions