從Numpy中的向量創建一個矩陣切片

Question

假設我有一個長度爲5的矩陣A，長度爲5的矢量b其元素指示矩陣A的對應行中需要多少個值。這意味着b中的每個值都以A的第二維的大小爲上界。我的問題是如何給定一個矩陣的一個片段給定一個向量，這是一個複雜的版本，通過編寫一個向量的整數值元素vector[:n]

例如，這可以通過循環遍歷A行：

import numpy 
A=numpy.arange(20).reshape((5,4)) 
b=numpy.array([0, 3, 3, 2, 3]) 
output=A[0, :b[0]] 
for i in xrange(1, A.shape[0]): 
    output=numpy.concatenate((output, A[i, :b[i]]), axis=0) 
# output is array([ 4, 5, 6, 8, 9, 10, 12, 13, 16, 17, 18]) 

當處理一個非常大的數組時，該循環的計算效率可能相當低。此外，我的目的是在沒有scan操作的情況下最終應用於Theano。我想避免使用循環來給出一個矢量的切片。

Answer 1

使用NumPy broadcasting的另一個好設置！

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

樣品運行

1）輸入：

In [16]: A 
Out[16]: 
array([[ 0, 1, 2, 3], 
     [ 4, 5, 6, 7], 
     [ 8, 9, 10, 11], 
     [12, 13, 14, 15], 
     [16, 17, 18, 19]]) 

In [17]: b 
Out[17]: array([0, 3, 3, 2, 3]) 

2）使用廣播來進行選擇創建掩模：

In [18]: b[:,None] > np.arange(A.shape[1]) 
Out[18]: 
array([[False, False, False, False], 
     [ True, True, True, False], 
     [ True, True, True, False], 
     [ True, True, False, False], 
     [ True, True, True, False]], dtype=bool) 

3）最後使用boolean-indexing用於選擇退出A：

In [19]: A[b[:,None] > np.arange(A.shape[1])] 
Out[19]: array([ 4, 5, 6, 8, 9, 10, 12, 13, 16, 17, 18]) 

Answer 2

你可以通過在列表中收集值，只是一個concatenate做加快循環：

In [126]: [A[i,:j] for i,j in enumerate(b)] 
Out[126]: 
[array([], dtype=int32), 
array([4, 5, 6]), 
array([ 8, 9, 10]), 
array([12, 13]), 
array([16, 17, 18])] 

In [127]: np.concatenate([A[i,:j] for i,j in enumerate(b)]) 
Out[127]: array([ 4, 5, 6, 8, 9, 10, 12, 13, 16, 17, 18])