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

【问题标题】：Creating a slice of a matrix from a vector in Numpy从 Numpy 中的向量创建矩阵切片
【发布时间】：2017-02-02 20:46:18
【问题描述】：

假设我有一个大小为 5 x 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 中应用它。我想避免使用循环在给定向量的情况下制作切片。

【问题讨论】：

标签： performance numpy indexing vectorization matrix-indexing

【解决方案1】：

您可以通过收集列表中的值来加速循环，然后只做一个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])

【讨论】：

【解决方案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])

【讨论】：