在numpy的2D数组中使用条件表达式指定适当的索引

7gs2gvoe 于 2023-08-05 发布在其他

关注(0)|答案(1)|浏览(110)

在2D数组的每一行中，如果第0个元素小于4，我想将10分配给第0和第2列中的元素。

import numpy as np
arr = np.array([[0,1,2],[3,4,5]])
arr[(arr[:,0] < 4), [0,2]] = 10
print(arr)

字符串
实际结果

[[10  1  2]
 [ 3  4 10]]

型
预期结果

[[10  1  10]
 [10  4  10]]

型

1条答案

将数组稍微大一点，使问题更有趣。

In [134]: arr = np.arange(9).reshape(3,3); arr
Out[134]: 
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

字符串
行选择，作为布尔值和索引：

In [135]: (arr[:,0] < 4)
Out[135]: array([ True,  True, False])    
In [136]: np.nonzero((arr[:,0] < 4))
Out[136]: (array([0, 1]),)

型
如果我们只是两个列表（或数组）（大小相同）;我们索引“对角线”。这就是你所得到的。

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

型
为了得到块或4个角，'行'索引必须是（2，1）形状：

In [138]: arr[[[0],[1]],[0,2]]
Out[138]: 
array([[0, 2],
       [3, 5]])

型
一个（2，1）数组broadcasts，其中有一个（2，）要选择为（2，2）块。
我们不能把布尔掩码扩展到二维。但是我们可以用np.ix_来创建等价的数组对：

In [140]: I,J = np.ix_((arr[:,0]<4), [0,2]); I,J
Out[140]: 
(array([[0],
        [1]]),
 array([[0, 2]]))

型
然后使用它们来索引arr：

In [141]: arr[I,J] = 10    
In [142]: arr
Out[142]: 
array([[10,  1, 10],
       [10,  4, 10],
       [ 6,  7,  8]])

型