如何使用高级布尔索引在 numpy 中赋值？

Question

我正在尝试使用布尔索引将一维输出广播到三维数组。我有一个要分配给的数组：

output_array = np.zeros((2,4,3))

然后是我用来做索引的一些布尔数组集：

dim0_bool = np.array([True, True])
dim0_dim1_bool = np.array([[True, True, True, False],
                           [False, True, True, True]])
dim0_dim2_bool = np.array([[True, True, False],
                           [False, True, True]])

从这些我可以构造一个三维布尔索引数组，并使用它为输出数组赋值：

output_array_idx = np.einsum('i, ij, ik -> ijk',
                            dim0_bool,
                            dim0_dim1_bool, 
                            dim0_dim2_bool)
output_array[output_array_idx] = 1.0

一切正常，并给出了预期的输出：

array([[[1., 1., 0.],
        [1., 1., 0.],
        [1., 1., 0.],
        [0., 0., 0.]],

       [[0., 0., 0.],
        [0., 1., 1.],
        [0., 1., 1.],
        [0., 1., 1.]]])

不过，我真正想做的（而且我不确定这是否可行）是将一维数组广播到 output_array 的索引元素。例如：

dim2_output = np.array([1.0, 2.0])

然后：

output_array[output_array_idx] = dim2_output

理想情况下会给出：

array([[[1., 2., 0.],
        [1., 2., 0.],
        [1., 2., 0.],
        [0., 0., 0.]],

       [[0., 0., 0.],
        [0., 1., 2.],
        [0., 1., 2.],
        [0., 1., 2.]]])

这不能按原样工作，因为通过 output_array_idx 索引 output_array 会将其缩减为一维数组，因此我无法分配 dim2_output（这是单一维度，但长度不同）。

希望这很清楚，如果有任何帮助（或关于如何更好地解决这个问题的建议），我们将不胜感激。

Answer 1

您可以构造一个 [1,2,1,2,...] 大小匹配的数组。

np.resize 进行正确的复制：

In [741]: np.resize([1,2],12)
Out[741]: array([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2])

或repeat:

In [744]: np.array([[1,2]]).repeat(6,axis=0).ravel()
Out[744]: array([1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2])

Answer 2

这是一个非常肮脏的方法。您可以将 output_array 设置为包含 dim2_output 的索引。缺点是 output_array 不再名副其实，您需要制作两份副本 - 一份用于索引，一份用于结果。

策略是初始化 nans 数组，根据您的布尔掩码填充数组，然后我们 np.nancumsum 恢复索引。

import numpy as np

# init output array with nans instead of zeros
output_array = np.full((2, 4, 3), np.nan)

dim0_bool = np.array([True, True])
dim0_dim1_bool = np.array([[True, True, True, False], [False, True, True, True]])
dim0_dim2_bool = np.array([[True, True, False], [False, True, True]])


output_array_idx = np.einsum(
    "i, ij, ik -> ijk", dim0_bool, dim0_dim1_bool, dim0_dim2_bool
)

# make the output array contain indices for dim2_output
output_array[output_array_idx] = 1
output_array = (np.nancumsum(output_array, -1) - 1) * output_array
output_array[np.isnan(output_array)] = -1

# append to work with index manipulations 0
dim2_output = np.array([1.0, 2.0, 0])

result = dim2_output[output_array.astype(np.int)]


# array([[[1., 2., 0.],
# [1., 2., 0.],
# [1., 2., 0.],
# [0., 0., 0.]],

# [[0., 0., 0.],
# [0., 1., 2.],
# [0., 1., 2.],
# [0., 1., 2.]]])

如何使用高级布尔索引在 numpy 中赋值？

How do I assign a value in numpy with advanced boolean indexing?

numpy

python-3.x

array-broadcasting