Numpy многомерный срез

Question

Если есть массив 2d Numpy:

array([[0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
       [0, 0, 0, 0, 1, 0, 1, 1, 1, 0],
       [1, 0, 1, 1, 1, 1, 1, 1, 1, 1],
       [0, 0, 1, 1, 1, 0, 1, 0, 0, 0],
       [0, 0, 1, 1, 1, 1, 1, 0, 1, 1]])

и я хотел нарезать каждую строку до первой позиции индекса, включая 1, включительно, как показано ниже:

array([[0, 1],
       [0, 0, 0, 0, 1],
       [1],
       [0, 0, 1,
       [0, 0, 1])

Можно ли добиться этого с помощью широковещания или все выходные массивы должны иметь одинаковую форму? У меня есть решение, использующее следующее, но мне было любопытно, можно ли это сделать с помощью радиовещания?

x = np.random.choice([0,1], size = [5,10])
idx = x.argmax(axis = 1) 

np.array([row[:i] for row, i in zip(x, idx + 1)])

Answer 1

Вы можете сделать это, используя dtype=object

a =np.array([[0, 1, 0, 0, 0, 1, 0, 0, 0, 1],
       [0, 0, 0, 0, 1, 0, 1, 1, 1, 0],
       [1, 0, 1, 1, 1, 1, 1, 1, 1, 1],
       [0, 0, 1, 1, 1, 0, 1, 0, 0, 0],
       [0, 0, 1, 1, 1, 1, 1, 0, 1, 1]]) 
idx = a.argmax(axis = 1) 
a = np.array([row[:i] for row, i in zip(a, idx + 1)], dtype=object)

Выход:

a = array([array([0, 1]), array([0, 0, 0, 0, 1]), array([1]),
       array([0, 0, 1]), array([0, 0, 1])], dtype=object)