How to efficiently unroll a matrix by value with numpy?

Broadcasted comparison is your friend:

B = (M[None, :] == np.arange(N)[:, None, None]).view(np.int8)

 np.array_equal(A, B)
# True

The idea is to expand the dimensions in such a way that the comparison can be broadcasted in the manner desired.


As pointed out by @Alex Riley in the comments, you can use np.equal.outer to avoid having to do the indexing stuff yourself,

B = np.equal.outer(np.arange(N), M).view(np.int8)

np.array_equal(A, B)
# True

You can make use of some broadcasting here:

P = np.arange(N)
Y = np.broadcast_to(P[:, None], M.shape)
T = np.equal(M, Y[:, None]).astype(int)

Alternative using indices:

X, Y = np.indices(M.shape)
Z = np.equal(M, X[:, None]).astype(int)

You can index into the identity matrix like so

 A = np.identity(N, int)[:, M]

or so

 A = np.identity(N, int)[M.T].T

Or use the new (v1.15.0) put_along_axis

A = np.zeros((N,5,5), int)
np.put_along_axis(A, M[None], 1, 0)

Note if N is much larger than 5 then creating an NxN identity matrix may be considered wasteful. We can mitigate this using stride tricks:

def read_only_identity(N, dtype=float):
    z = np.zeros(2*N-1, dtype)
    s, = z.strides
    z[N-1] = 1
    return np.lib.stride_tricks.as_strided(z[N-1:], (N, N), (-s, s))