如何在具有给定起点索引的numpy数组中选择元素

时间:2018-07-20 11:15:50

标签: python numpy

例如,我有一个像这样的矩阵:

In [2]: a = np.arange(12).reshape(3, 4)

In [3]: a
Out[3]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

和起点索引数组:

In [4]: idx = np.array([1, 2, 0])

In [5]: idx
Out[5]: array([1, 2, 0])

是否有任何矢量化的方式来做这种事情:

for i in range(3):
    # The following are some usecases
    a[i, idx[i]:] = 0
    a[i, idx[i]-1:] = 0
    a[i, :idx[i]] = 0
    a[i, idx[i]:idx[i]+2] = 0

编辑:预期输出:

array([[ 0,  x,  x,  x],
       [ 4,  5,  x,  x],
       [ x,  x,  x,  x]])

x是占位符,指示我要选择的内容。

2 个答案:

答案 0 :(得分:0)

未提供预期的准确输出。因此,我认为以下内容可能对您大有帮助。

>>> a = np.arange(12).reshape(3, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> row = np.array([0, 1, 2])
>>> col = np.array([1, 2, 0])
>>> a[row, col]
array([1, 6, 8])

您可以将row的{​​{1}}和col s设置为一个值:

a

答案 1 :(得分:0)

aproach也适用于矩形矩阵。创建布尔面罩槽广播:

a = np.arange(12).reshape(3, 4)
idx = np.array([1, 2, 0])
mask=np.arange(a.shape[1]) >= idx[:,None]
mask
#array([[False,  True,  True,  True],
#       [False, False,  True,  True],
#       [ True,  True,  True,  True]], dtype=bool)

例如,使用您的占位符 -1,并设置a的值,其中mask等于该占位符:

x = -1
a[mask] = x
a
#array([[ 0, -1, -1, -1],
#       [ 4,  5, -1, -1],
#      [-1, -1, -1, -1]])