numpy select each row with different lengths

Question

I have two array s

x=array([[0, 0, 0, 0, 0],
    [1, 0, 0, 0, 0],
    [2, 2, 2, 2, 2]])

I want to subselect elements in each row by the length in array y

y = array([3, 2, 4])

My target is z:

z = array([[0, 0, 0],
   [1, 0,],
   [2, 2, 2, 2]])

How could I do that with numpy functions instead of list/loop?

Thank you so much for your help.

Answer 1

Numpy array is optimized for homogeneous array with a specific dimensions. I like to think of it like a matrix: it does not make sense to have a matrix with different number of elements on each rows.

That said, depending on how you want to use the processed array, you can simply make a list of array:

z = [array([0, 0, 0]),
array([1, 0,]),
array([2, 2, 2, 2]])]

Still, you will need to do that manually:

x = array([[0, 0, 0, 0, 0], [1, 0, 0, 0, 0], [2, 2, 2, 2, 2]])
y = array([3, 2, 4])

z = [x_item[:y_item] for x_item, y_item in zip(x, y)]

The list comprehension iterates over the x and y combined with zip() to create the new slice of the original array.

Answer 2

Something like this also,

z = [x[i,:e] for i,e in enumerate(y)]