Индекс Numpy: первое (различное) количество элементов из каждой строки в массиве 2d

Question

(короткая версия моего вопроса: в numpy, есть ли элегантный способ эмулировать tf.sequence_mask из тензорного потока?)

У меня есть двумерный массив a (каждая строка представляет последовательность различной длины). Далее идет 1d массив b (представляющий длины последовательности). Существует ли элегантный способ получить (уплощенный) массив, который будет содержать только такие элементы a, которые относятся к последовательностям, указанным их длиной b:

a = np.array([
    [1, 2, 3, 2, 1],  # I want just [:3] from this row
    [4, 5, 5, 5, 1],  # [:2] from this row
    [6, 7, 8, 9, 0]   # [:4] from this row
])
b = np.array([3,2,4])  # 3 elements from the 1st row, 2 from the 2nd, 4 from the 4th row

желаемый результат:

[1, 2, 3, 4, 5, 6, 7, 8, 9]

Под elegant way я имею в виду то, что избегает петель.

Answer 1

Используйте broadcasting, чтобы создать маску той же формы, что и 2D-массив, а затем просто замаскируйте и извлеките действительные элементы -

a[b[:,None] > np.arange(a.shape[1])]

Пробный прогон -

In [360]: a
Out[360]: 
array([[1, 2, 3, 2, 1],
       [4, 5, 5, 5, 1],
       [6, 7, 8, 9, 0]])

In [361]: b
Out[361]: array([3, 2, 4])

In [362]: a[b[:,None] > np.arange(a.shape[1])]
Out[362]: array([1, 2, 3, 4, 5, 6, 7, 8, 9])