numpy使用切片修改多个值的二维数组

时间:2019-01-13 22:34:09

标签: python arrays numpy slice submatrix

我想根据另一个数组的值更改numpy 2D数组中的某些值。使用布尔切片选择子矩阵的行,使用整数切片选择列。

这是一些示例代码:

import numpy as np

a = np.array([
    [0, 0, 1, 0, 0],
    [1, 1, 1, 0, 1],
    [0, 1, 0, 1, 0],
    [1, 1, 1, 0, 0],
    [1, 0, 0, 0, 1],
    [0, 0, 0, 0, 0],
])

b = np.ones(a.shape)    # Fill with ones
rows = a[:, 3] == 0     # Select all the rows where the value at the 4th column equals 0
cols = [2, 3, 4]        # Select the columns 2, 3 and 4

b[rows, cols] = 2       # Replace the values with 2
print(b)

我想要在b中得到的结果是:

[[1. 1. 2. 2. 2.]
 [1. 1. 2. 2. 2.]
 [1. 1. 1. 1. 1.]
 [1. 1. 2. 2. 2.]
 [1. 1. 2. 2. 2.]
 [1. 1. 2. 2. 2.]]

但是,我唯一得到的是一个例外:

IndexError
shape mismatch: indexing arrays could not be broadcast together with shapes (5,) (3,)

我如何获得想要的结果?

1 个答案:

答案 0 :(得分:1)

您可以使用argwhere

rows = np.argwhere(a[:, 3] == 0)    
cols = [2, 3, 4]        

b[rows, cols] = 2       # Replace the values with 2
print(b)

输出

[[1. 1. 2. 2. 2.]
 [1. 1. 2. 2. 2.]
 [1. 1. 1. 1. 1.]
 [1. 1. 2. 2. 2.]
 [1. 1. 2. 2. 2.]
 [1. 1. 2. 2. 2.]]