numpy: insert array in place of an element

Question

Input:

import numpy as np

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

Desired output:

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

What I've tried:

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

The above code gives me a ValueError: operands could not be broadcast together with shapes (<a's dimentions>) (2,) (2,)

Edit: Bad example on my end. I thought making it simple would make it more readable. But I don't necessarily want ones in place of 1 and zeros in place of 0, and the input isn't limited to 0s and 1s. Instead, I want the input to be something like:

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

and the output to be:

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

Answer 1

You can expand dimensions and then repeat along the new dimension(or repeat then reshape):

np.repeat(np.expand_dims(a, axis=2), 2, axis=2)

Output:

[[[0 0]
  [1 1]
  [0 0]]

 [[1 1]
  [0 0]
  [0 0]]]

And if you want to have different output, now you can apply your condition to new array easier.

UPDATE: per edit on the post for a general case: Assuming you have a mapping dictionary (depending on your mapping, solution could be different)

map_dict = {0: [5, 6],
     1: [8, 2],
     2: [7, 4]}

indexer = np.array([map_dict.get(i, [-1,-1]) for i in range(np.amax(a)+1)])[a]

or in a more specific case if your mapping includes all integers between 0 and max(a), like your example, use this:

indexer = np.array([map_dict[i] for i in np.unique(a)])[a]

output:

[[[5 6]
  [8 2]
  [5 6]]

 [[8 2]
  [7 4]
  [5 6]]]