从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])