python-3.x - 如何使用高级 bool 索引在 numpy 中分配值？

Question

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

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

然后是我用来进行索引的一些 bool 数组集:

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

根据这些，我可以构造一个三维 bool 索引数组，并使用它为输出数组分配一个值:

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

或重复:

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