Performance 从Numpy中的向量创建矩阵切片

假设我有一个大小为5×4的矩阵

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

---

在处理非常大的阵列时，此循环的计算效率可能相当低。此外，我的目的是最终在无需

扫描操作的情况下将其应用于Theano中。我想避免使用循环来生成给定向量的切片。
另一个很好的使用设置

样本运行

1） 投入：

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） 最后使用
布尔索引
选择元素关闭
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])

您可以通过在列表中收集值，并只执行一次连接来加速循环：

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

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

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