Constructing A Special Matrix In Numpy Dynamically

So my objective is the following, given the size of the matrix s, I am attempting to create a matrix that looks like the following, but for size sxs: [1 1 0] [1 1 1] [0 1 1] For t

Solution 1:

Method #1: Instead of a bunch of 2x2 matrices, it might be easier to look at it as three diagonals of 1 and combine those:

>>> s = 3
>>> np.diag([1]*s,0) + np.diag([1]*(s-1),-1) + np.diag([1]*(s-1), 1)
array([[1, 1, 0],
       [1, 1, 1],
       [0, 1, 1]])
>>> s = 4
>>> np.diag([1]*s,0) + np.diag([1]*(s-1),-1) + np.diag([1]*(s-1), 1)
array([[1, 1, 0, 0],
       [1, 1, 1, 0],
       [0, 1, 1, 1],
       [0, 0, 1, 1]])

Method #2: (inspired by Divankar's answer), we can think in terms of distance from the centre:

>>>s = 4>>>i,j = np.indices((s,s))>>>(abs(i-j) <= 1).astype(int)
array([[1, 1, 0, 0],
       [1, 1, 1, 0],
       [0, 1, 1, 1],
       [0, 0, 1, 1]])

Method #3: we could take advantage of tril or triu and do some arithmetic:

>>> m = np.tril(np.ones((s,s)),1)
>>> m * m.T
array([[ 1.,  1.,  0.,  0.,  0.],
       [ 1.,  1.,  1.,  0.,  0.],
       [ 0.,  1.,  1.,  1.,  0.],
       [ 0.,  0.,  1.,  1.,  1.],
       [ 0.,  0.,  0.,  1.,  1.]])
>>> m = np.tril(np.ones((s,s)),2)
>>> m * m.T
array([[ 1.,  1.,  1.,  0.,  0.],
       [ 1.,  1.,  1.,  1.,  0.],
       [ 1.,  1.,  1.,  1.,  1.],
       [ 0.,  1.,  1.,  1.,  1.],
       [ 0.,  0.,  1.,  1.,  1.]])

Solution 2:

Vectorized approach with broadcasting -

A = np.arange(s)
out = ((A[:,None] < A+2) & (A[:,None] > A-2)).astype(int)

Sample run -

In [60]: s =3
    ...: A = np.arange(s)
    ...: out= ((A[:,None] < A+2) & (A[:,None] > A-2)).astype(int)
    ...: 

In [61]: outOut[61]: 
array([[1, 1, 0],
       [1, 1, 1],
       [0, 1, 1]])

In [62]: s =4
    ...: A = np.arange(s)
    ...: out= ((A[:,None] < A+2) & (A[:,None] > A-2)).astype(int)
    ...: 

In [63]: outOut[63]: 
array([[1, 1, 0, 0],
       [1, 1, 1, 0],
       [0, 1, 1, 1],
       [0, 0, 1, 1]])

Solution 3:

You could use sympy.Matrix also:

from sympy import Matrix
Matrix(4, 4, lambda i,j: 1if (-2<i-j<2) else0)

Matrix([
[1, 1, 0, 0],
[1, 1, 1, 0],
[0, 1, 1, 1],
[0, 0, 1, 1]])