Insert elements in between in numpy array after a certain value-CodePudding

I have a numpy array:

data=np.array([0,0,0,0,0,0,1,0,0,1,0,0,1,0,0,1,0,0,0,0,0,0,1,0,0,1,0,0,1,0,0,1,0,0,0,0,0,0]).

And I want to fill in the gaps of zeroes between ones, where ones are separated with no more than 2 zeroes. The output should be:

data=np.array([0,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,0,1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,0]).

I write the code and it apparently does the job, but I get an error.

    if (data[i]==1) & (data[i 3]==1):
      for j in range(3):
        data[i j]=1
       if data[-1]==1:   #If 1 is at the end of the array it should not do anything, but I guess it does not work.
         pass```

IndexError: index 38 is out of bounds for axis 0 with size 38.

Could you help me to improve?

CodePudding user response：

It's possible to avoid using loops by shifting the values in data instead of looping. The code below copes with gaps of 1 or 2 zeroes and groups of consecutive 1s.

import numpy as np 

def fill_zeros( data ):
    result = data.copy()
    temp = np.zeros_like( data )
    temp[2:] = data[2:] & data[:-2]   # Gaps of one.
    result[ 1:-1 ] |= temp[2: ]       # temp shifted left one.
    temp[:] = 0                       # reset temp to zeroes
    temp[3:] = data[3:] & data[:-3]   # Gaps of two
    result[1:-2] |= temp[ 3:]         # Shift temp 2 left
    result[2:-1] |= temp[ 3:]         # Shift temp 1 left
    return result

To test:

data=np.array([0,0,0,0,0,0,1,0,0,0,1,0,1,0,0,0,1,0,0,0,0,0,0,1,0,0,1,1,0,0,1,0,0,1,0,0,0,0,0])
data
#  array([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,
#      0, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0])

fill_zeros( data )
# array([0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0,
#      0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0])

CodePudding user response：

assuming that there are not any two consecutive 1's

# find index of first 1's
idx = np.where(np.logical_and(data[:-2] == 1, data[2:] == 1))  
idx = idx[0][::2]

# set two elements after first 1's to 1
data[idx 1] = 1
data[idx 2] = 1