Create an NxM matrix A to an NxMxL matrix B where B[i,j,:] = kronecker_delta(A[i,j])?

Question

Is there a way to convert a NxM matrix A where

• all values of A are positive integers

to an NxMxL matrix B where

• L = 1 + max(A)
• B[i,j,k] = {1 if k==A[i,j] and 0 otherwise}

using loops I have done the following:

B = np.zeros((A.shape[0],A.shape[1],1+np.amax(A)))
for i in range(A.shape[0]):
    for j in range(A.shape[1]):
        B[i,j,A[i,j]] = 1

the solution ideally avoids any use of for loops and uses only slicing or indexing or numpy functions

Answer 1

A sample A:

In [231]: A = np.array([1,0,3,2,2,4]).reshape(2,3)
In [232]: A
Out[232]: 
array([[1, 0, 3],
       [2, 2, 4]])

Your code and B:

In [233]: B = np.zeros((A.shape[0],A.shape[1],1+np.amax(A)))
     ...: for i in range(A.shape[0]):
     ...:     for j in range(A.shape[1]):
     ...:         B[i,j,A[i,j]] = 1
     ...:             
In [234]: B
Out[234]: 
array([[[0., 1., 0., 0., 0.],
        [1., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0.]],

       [[0., 0., 1., 0., 0.],
        [0., 0., 1., 0., 0.],
        [0., 0., 0., 0., 1.]]])

Defining arrays to index B in way that broadcasts with A:

In [235]: I,J=np.ix_(np.arange(2),np.arange(3))

In [236]: B[I,J,A]
Out[236]: 
array([[1., 1., 1.],
       [1., 1., 1.]])

Use that indexing to change all the 1s of B to 20:

In [237]: B[I,J,A]=20

In [238]: B
Out[238]: 
array([[[ 0., 20.,  0.,  0.,  0.],
        [20.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0., 20.,  0.]],

       [[ 0.,  0., 20.,  0.,  0.],
        [ 0.,  0., 20.,  0.,  0.],
        [ 0.,  0.,  0.,  0., 20.]]])

the indexes (2,1) and (1,3) pair with (2,3):

In [239]: I,J
Out[239]: 
(array([[0],
        [1]]),
 array([[0, 1, 2]]))

There's also newer pair of functions that do that same thing. I'm more familiar with the earlier method

In [241]: np.take_along_axis(B,A[:,:,None],2)
Out[241]: 
array([[[20.],
        [20.],
        [20.]],

       [[20.],
        [20.],
        [20.]]])
    
In [243]: np.put_along_axis(B,A[:,:,None],1,axis=2)

In [244]: B
Out[244]: 
array([[[0., 1., 0., 0., 0.],
        [1., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0.]],

       [[0., 0., 1., 0., 0.],
        [0., 0., 1., 0., 0.],
        [0., 0., 0., 0., 1.]]])