python - 高级布尔索引

Question

我想通过掩码选择值并使用掩码数组更改值。

代码：

import numpy as np

a = np.zeros((2, 2), dtype=(np.uint8, 3))
x = np.arange(4, dtype=int).reshape((2, 2))

mask = np.logical_and(a1 < 3, a1 > 0)

a[mask] = (1, x[mask], 2)

我想要结果：

a[mask]
>> [[1, 1, 2], [1, 2, 2]]

但我得到错误： ValueError: setting an array element with a sequence.

a[mask] = (1, 2, 2) 如果尝试做数组之类的事情

[[[0, 0, 0],
[1, 2, 2]],
[[1, 2, 2],
[0, 0, 0]]]

但我需要使用 x 的值。让它看起来像

[[[0, 0, 0],
[1, 1, 3]],
[[1, 2, 3],
[0, 0, 0]]]

我该怎么做？

Answer 1

可以分两步完成。

import numpy as np 

a = np.zeros((2, 2), dtype=(np.uint8, 3)) 
x = np.arange(4, dtype=int).reshape((2, 2)) 
a1 = x  # Create an a1 for the mask

mask = np.logical_and(a1 < 3, a1 > 0) 

a[mask] = (1, 0, 2)      # Set the outer columns                                         
a[mask, 1] = x[mask]     # Set the column 1

print( a )                                                              
# [[[0 0 0]
#   [1 1 2]]

#  [[1 2 2]
#  [0 0 0]]]