修改2D NumPy数组的每一行中的不同列

时间:2017-05-01 21:19:52

标签: python arrays numpy indexing

我有以下问题:

让我们说我有一个像这样定义的数组:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])

我想要做的是使用Numpy多重索引并将几​​个元素设置为0.为此,我创建了一个向量:

indices_to_remove = [1, 2, 0]

我想要它的意思如下:

  1. 删除索引为' 1'的元素从第一行
  2. 删除索引' 2'从第二行
  3. 删除索引' 0'从第三行
  4. 结果应该是数组[[1,0,3],[4,5,0],[0,8,9]]

    我已设法通过以下代码获取我想要修改的元素的值:

    values = np.diagonal(np.take(A, indices, axis=1))
    

    然而,这并不允许我修改它们。怎么能解决这个问题?

1 个答案:

答案 0 :(得分:2)

您可以使用integer array indexing分配这些零 -

A[np.arange(len(indices_to_remove)), indices_to_remove] = 0

示例运行 -

In [445]: A
Out[445]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

In [446]: indices_to_remove
Out[446]: [1, 2, 0]

In [447]: A[np.arange(len(indices_to_remove)), indices_to_remove] = 0

In [448]: A
Out[448]: 
array([[1, 0, 3],
       [4, 5, 0],
       [0, 8, 9]])