I have a matrix M
with values 0 through N
within it. I'd like to unroll this matrix to create a new matrix A
where each submatrix A[i, :, :]
represents whether or not M == i.
The solution below uses a loop.
# Example Setup
import numpy as npnp.random.seed(0)
N = 5
M = np.random.randint(0, N, size=(5,5))# Solution with Loop
A = np.zeros((N, M.shape[0], M.shape[1]))
for i in range(N):A[i, :, :] = M == i
This yields:
M
array([[4, 0, 3, 3, 3],[1, 3, 2, 4, 0],[0, 4, 2, 1, 0],[1, 1, 0, 1, 4],[3, 0, 3, 0, 2]])M.shape
# (5, 5)A
array([[[0, 1, 0, 0, 0],[0, 0, 0, 0, 1],[1, 0, 0, 0, 1],[0, 0, 1, 0, 0],[0, 1, 0, 1, 0]],...[[1, 0, 0, 0, 0],[0, 0, 0, 1, 0],[0, 1, 0, 0, 0],[0, 0, 0, 0, 1],[0, 0, 0, 0, 0]]])A.shape
# (5, 5, 5)
Is there a faster way, or a way to do it in a single numpy operation?