how to slice 2D numpy array based on 2 1D arrays containing initial and final indexes

Question

I have a 2D numpy array, let's say it has shape 4x10 (4 rows and 10 columns). I have 2 1D arrays that have the initial and final indexes, so they are both 20x1. For an example, let's say

initial = [1, 2, 4, 5]
final = [3, 6, 8, 6]

then I'd like to get

data[0,1:3]
data[1,2:6]
data[2,4:8]
data[3,5:6]

Of course, each of those arrays would have different size, so I'd like to store them in a list. If I were to do it with a for loop, it'd look like this:

arrays = []
for i in range(4):
    slice = data[i,initial[i]:final[i]]
    arrays.append(slice)

Is there a more efficient way to do this? I'd rather avoid use a for loop, because my actual data is huge.

4x10 (10 rows and 10 columns) ? Maybe there's some issue in the numbers you mention.. — Eskapp
– Eskapp, Commented Feb 7, 2021 at 1:33
How else are you going to get a list? You could write the loop as a comprehension. If the slices had the same length there's a chance of constructing a 2d array. But not with differinvg lengths. — hpaulj
– hpaulj, Commented Feb 7, 2021 at 2:04
it doesn't have to be a list, but I don't know how to store arrays of different lengths, other than using a list. — Schach21
– Schach21, Commented Feb 7, 2021 at 2:43

Pablo C · Accepted Answer · 2021-02-07 02:54:41Z

1

You can use numpy.split with flattened data (using numpy.ndarray.flatten) and modifying the slices:

sections = np.column_stack([initial, final]).flatten()

sections[::2] += np.arange(len(initial)) * data.shape[1]
sections[1::2] += sections[::2] - np.array(initial)

np.split(data.flatten(), sections)[1::2]

answered Feb 7, 2021 at 2:54

Pablo C

4,7612 gold badges10 silver badges26 bronze badges

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

Collectives™ on Stack Overflow

how to slice 2D numpy array based on 2 1D arrays containing initial and final indexes

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related