高级布尔索引

问题描述

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

代码：

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) 之类的事情 数组将是

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

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

[[[0,3]],3],0]]]

我该怎么做？

解决方法

可以分两步完成。

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,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]]]