从具有索引的二维矩阵构建 3D boolean 矩阵，在 numpy

Question

I have a 2D matrix of shape (3, 4) with indices ranging from 0 to 8:

a = array([[0, 4, 1, 2],
           [5, 0, 2, 3],
           [8, 6, 0, 5]])

Currently, I use a for loop to build a 3D boolean array of shape (9, 3, 4) that stores True at the position of each index, for each row between 0 and 8:

b = np.zeros((9, 3, 4), dtype=bool)
for i in range(9):
    b[i] = np.where(a == i, True, False)

Is there a way to achieve the same result without iteration, perhaps using numpy functions?

Answer 1

Is this something you're looking for?

import numpy as np

a = np.array([[0, 4, 1, 2],
              [5, 0, 2, 3],
              [8, 6, 0, 5]])
y, x = np.mgrid[0:a.shape[0], 0:a.shape[1]]

data = np.zeros((9,) + a.shape, dtype=bool)
data[a, y, x] = True

Answer 2

A very short solution that takes advantage of numpy broadcasting:

b = np.array([a]*9) == np.arange(9).reshape(-1,1,1)

Output:

>>> b
array([[[ True, False, False, False],
        [False,  True, False, False],
        [False, False,  True, False]],

       [[False, False,  True, False],
        [False, False, False, False],
        [False, False, False, False]],

       [[False, False, False,  True],
        [False, False,  True, False],
        [False, False, False, False]],

       [[False, False, False, False],
        [False, False, False,  True],
        [False, False, False, False]],

       [[False,  True, False, False],
        [False, False, False, False],
        [False, False, False, False]],

       [[False, False, False, False],
        [ True, False, False, False],
        [False, False, False,  True]],

       [[False, False, False, False],
        [False, False, False, False],
        [False,  True, False, False]],

       [[False, False, False, False],
        [False, False, False, False],
        [False, False, False, False]],

       [[False, False, False, False],
        [False, False, False, False],
        [False, False, False, False]]])