numpy：在元素处插入数组

Question

输入：

import numpy as np

a = np.array([
    [0, 1, 0],
    [1, 0, 0]
])

所需的输出：

b = np.array([
    [[0, 0],[1, 1], [0, 0]],
    [[1, 1], [0, 0], [0, 0]
])

我尝试过的：

b = np.where(a == 0, np.array([0, 0]), np.array([1, 1]))

上面的代码给了我ValueError: operands could not be broadcast together with shapes (<a's dimentions>) (2,) (2,)

Edit：我这方面的例子很糟糕。我认为简化它会使它更具可读性，但是我不一定要用ones代替1，用zeros代替0。相反，我希望输出是这样的：

b = np.array([
    [[5, 6],[8, 2], [5, 6]],
    [[8, 2], [5, 6], [5, 6]
])

Answer 1

您可以扩展尺寸，然后沿新尺寸重复（或重复然后重新塑形）：

Answer 2

您可以使用遮罩和numpy broadcasting