What is the fastest way to insert elements diagonally in 2D numpy array?

Question

Suppose we have a 2D numpy array like:

matrix = [[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9],
          [10, 11, 12]]

I want to insert a value say 0 diagonally such that it becomes:

matrix = [[0, 1, 2, 3],
          [4, 0, 5, 6],
          [7, 8, 0, 9],
          [10, 11, 12, 0]]

What is the fastest way to do that?

Answer 1

Create a new bigger matrix, that have space left for the zeros. Copy the original matrix to a submatrix, clip and reshape:

matrix = numpy.array([[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9],
          [10, 11, 12]])

matrix_new = numpy.zeros((4,5))
matrix_new[:-1,1:] = matrix.reshape(3,4)
matrix_new = matrix_new.reshape(-1)[:-4].reshape(4,4)

or in a more generalized form:

matrix = numpy.array([[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9],
          [10, 11, 12]])

d = matrix.shape[0]
assert matrix.shape[1] == d - 1
matrix_new = numpy.ndarray((d, d+1), dtype=matrix.dtype)
matrix_new[:,0] = 0
matrix_new[:-1,1:] = matrix.reshape((d-1, d))
matrix_new = matrix_new.reshape(-1)[:-d].reshape(d,d)

Answer 2

Here's one way (but I can't promise that it is the fastest way):

In [62]: a
Out[62]: 
array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

In [63]: b = np.zeros((a.shape[0], a.shape[1]+1), dtype=a.dtype)

In [64]: i = np.arange(b.shape[0])

In [65]: j = np.arange(b.shape[1])

In [66]: b[np.not_equal.outer(i, j)] = a.ravel()  # or a.flat, if a is C-contiguous

In [67]: b
Out[67]: 
array([[ 0,  1,  2,  3],
       [ 4,  0,  5,  6],
       [ 7,  8,  0,  9],
       [10, 11, 12,  0]])

It works for any 2-d array a :

In [72]: a
Out[72]: 
array([[17, 18, 15, 19, 12],
       [16, 14, 11, 16, 17],
       [19, 11, 16, 11, 14]])

In [73]: b = np.zeros((a.shape[0], a.shape[1]+1), dtype=a.dtype)

In [74]: i = np.arange(b.shape[0])

In [75]: j = np.arange(b.shape[1])

In [76]: b[np.not_equal.outer(i, j)] = a.flat

In [77]: b
Out[77]: 
array([[ 0, 17, 18, 15, 19, 12],
       [16,  0, 14, 11, 16, 17],
       [19, 11,  0, 16, 11, 14]])

It works, but I think @Daniel's answer is the way to go.

Answer 3

another approach, probably slower, with append and reshape

import numpy as np

mat = np.array(range(1,13)).reshape(4,3)
mat

array([[ 1,  2,  3],
       [ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

z=np.zeros((3,1), dtype=mat.dtype)
m3=np.append(z,mat.reshape(3,4),1)
np.append(m3,0).reshape(4,4)

array([[ 0,  1,  2,  3],
       [ 4,  0,  5,  6],
       [ 7,  8,  0,  9],
       [10, 11, 12,  0]])

Answer 4

Looks like you are taking the lower and upper triangular arrays and separating them with a diagonal of zeros. This sequence does that:

In [54]: A=np.arange(1,13).reshape(4,3)

Target array, with one more column

In [55]: B=np.zeros((A.shape[0],A.shape[1]+1),dtype=A.dtype)

Copy over the lower tri (without the diagonal)

In [56]: B[:,:-1]+=np.tril(A,-1)

Copy the upper tri

In [57]: B[:,1:]+=np.triu(A,0)

In [58]: B
Out[58]: 
array([[ 0,  1,  2,  3],
       [ 4,  0,  5,  6],
       [ 7,  8,  0,  9],
       [10, 11, 12,  0]])

There are some np.tril_indices... functions, but they only work with square arrays. So they can't be used with A .

Answer 5

Lets say you have apxq numpy 2d array A, here is a sample with (p,q) as (3,4):

In []: A = np.arange(1,13).reshape(4,3)
In []: A
Out[]: 
array([[ 1,  2,  3],
      [ 4,  5,  6],
      [ 7,  8,  9],
      [10, 11, 12]])

Step 1:

To insert a diagonal of zeros, it will require making a new 2d array of shape px q+1.

Before this we create a 2d array with column index values of non-diagonal elements for the new 2d array like this

In []: columnIndexArray = np.delete(np.meshgrid(np.arange(q+1), np.arange(p))[0], np.arange(0, p * (q+1), q+2)).reshape(p,q)

The output of the above will look as follows:

In []: columnIndexArray
Out[]: 
array([[1, 2, 3],
      [0, 2, 3],
      [0, 1, 3],
      [0, 1, 2]])

Step 2:

Now construct px q+1 2d array of zeros like this

In []: B = np.zeros((p,q+1))

In []: B
Out[]: 
array([[ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.]])

Step 3:

Now assign the non diagonal elements with the values from A

In []: B[np.arange(p)[:,None], columnIndexArray] = A

In []: B
Out[]: 
array([[  0.,   1.,   2.,   3.],
      [  4.,   0.,   5.,   6.],
      [  7.,   8.,   0.,   9.],
      [ 10.,  11.,  12.,   0.]])

Note: To make your code dynamic replace p with A.shape[0] and q with A.shape[1] respectively.

Answer 6

Again, I don't know how fast it is but you could try using numpy.lib.stride_tricks.as_strided :

import numpy as np
as_strided = np.lib.stride_tricks.as_strided

matrix = (np.arange(12)+1).reshape((4,3))

n, m = matrix.shape
t = matrix.reshape((m, n))
t = np.hstack((np.array([[0]*m]).T, t))
t = np.vstack((t, [0]*(n+1)))
q = as_strided(t, (n,n), (t.itemsize*n, 8))
print(q)

Output:

[[ 0  1  2  3]
 [ 4  0  5  6]
 [ 7  8  0  9]
 [10 11 12  0]]

That is, pad the reshaped array left and bottom with zeros and stride to put the left-hand zeros on the diagonal. Unfortunately, you need the bottom row of zeros to get the final zero in your output matrix in the case of (n+1,n) arrays (since eg 5*3=15 is one less than 4*4=16). You can do without this if you start with a square matrix.