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

时间:2016-09-25 12:52:46

标签: performance numpy indexing vectorization matrix-indexing

如果我有一个大小为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])

在处理非常大的数组时,此循环的计算效率可能相当低。此外,我的目的是最终在没有scan操作的情况下将其应用于Theano。我想避免使用循环来给出一个给定矢量的切片。

2 个答案:

答案 0 :(得分:2)

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

答案 1 :(得分:0)

您可以通过收集列表中的值并仅执行一个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])