Ayush Chaurasia
Ayush Chaurasia

Reputation: 647

Insert values in alternate rows and columns in numpy arrays

I need to insert arbitrary number of zeros to alternate rows and columns of a numpy array. For example, lets suppose that we want to insert 1 zero in all alternate columns and rows.

Input => [[ 1,2,3],
           [ 4,5,6],
           [ 7,8,9]]

output => [[ 1,0,2,0,3],
           [ 0,0,0,0,0],
           [ 4,0,5,0,6],
           [ 0,0,0,0,0],
           [ 7,0,8,0,9]] 

I know how this can be achieved by using loops but I'm not sure if that's the most efficient way or there is some vectorized implementation possible .

Upvotes: 3

Views: 2284

Answers (1)

Divakar
Divakar

Reputation: 221574

One approach with zeros-initialization and using step-sized slicing for assigning -

def insert_zeros(a, N=1):
    # a : Input array
    # N : number of zeros to be inserted between consecutive rows and cols 
    out = np.zeros( (N+1)*np.array(a.shape)-N,dtype=a.dtype)
    out[::N+1,::N+1] = a
    return out

Sample runs -

In [167]: a
Out[167]: 
array([[83, 87, 19, 24],
       [28, 24, 24, 77],
       [26, 87, 57, 37]])

In [168]: insert_zeros(a, N=1)
Out[168]: 
array([[83,  0, 87,  0, 19,  0, 24],
       [ 0,  0,  0,  0,  0,  0,  0],
       [28,  0, 24,  0, 24,  0, 77],
       [ 0,  0,  0,  0,  0,  0,  0],
       [26,  0, 87,  0, 57,  0, 37]])

In [169]: insert_zeros(a, N=2)
Out[169]: 
array([[83,  0,  0, 87,  0,  0, 19,  0,  0, 24],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [28,  0,  0, 24,  0,  0, 24,  0,  0, 77],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [26,  0,  0, 87,  0,  0, 57,  0,  0, 37]])

Upvotes: 9

Related Questions