基于 numpy 1D 数组的 2D numpy 数组中的智能分配

Question

我有一个 numpy 2D 数组，我想根据以下逻辑将其转换为 -1\\1 值：

一种。 找到每一行的 argmax()

湾基于该一维数组 (a) 分配包含值 1 的值

C。 基于此一维数组的否定赋值 -1

例子：

arr2D = np.random.randint(10,size=(3,3))
idx = np.argmax(arr2D, axis=1)

arr2D = [[5 4 1]
         [0 9 4]
         [4 2 6]]
idx = [0 1 2]

arr2D[idx] = 1
arr2D[~idx] = -1

我得到的是这样的：

arr2D = [[-1 -1 -1]
         [-1 -1 -1]
         [-1 -1 -1]]

虽然我想要：

arr2D = [[1 -1 -1]
         [-1 1 -1]
         [-1 -1 1]]

感谢一些帮助，谢谢

Answer 1

方法#1

用那些argmax创建一个掩码 -

mask = idx[:,None] == np.arange(arr2D.shape[1])

然后，使用这些索引，然后使用它来创建那些 1s 和 -1s 数组 -

out = 2*mask-1

或者，我们可以使用np.where -

out = np.where(mask,1,-1)

方法#2

创建面具的另一种方法是 -

mask = np.zeros(arr2D.shape, dtype=bool)
mask[np.arange(len(idx)),idx] = 1

然后，使用方法#1 中列出的方法之一out 。

方法#3

另一种方式是这样 -

out = np.full(arr2D.shape, -1)
out[np.arange(len(idx)),idx] = 1

或者，我们可以使用np.put_along_axis进行分配 -

np.put_along_axis(out,idx[:,None],1,axis=1)