Numpy索引：第一個（變化的）2d數組中每行的元素數

Question

（我的問題的簡短版本：在numpy中，是否有一種從tensorflow模擬tf.sequence_mask的優雅方法？）

我有一個2d數組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])