In a one dimensional numpy array of 1's and 0's, how can I convert the next n elements following a 1 to 0's?

user_15199888 0
Viewed 131

For a one dimensional numpy array of 1's and 0's, how can I effectively "mask" the array such that after the occurrence of a 1, the next n elements of the array are converted to zero. After the n elements have passed, the pattern repeats such that the first next occurrence of a 1 is preserved followed again by n zeros.

It is important that the first eligible occurrences of 1 are preserved, so a simple mask such as: [true, false, false, true ...] won't work. furthermore, the data set is massive so efficiency is important.

I've written crude python code to give me the desired results, but it is way too slow for what I need. Here is an example:

data = [0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1]
n = 3

newData = []
tail = 0
for x in data:
    if x == 1 and tail <= 0:
        newData.append(1)
        tail = n
    else:
        newData.append(0)
        tail -= 1
print(newData)

newData: [0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1]

Is there possibly a vectorized numpy solution to this problem? I'm processing tens of thousands of arrays, with more than a million elements in each array. So far using numpy functions has been the only way to manage this.

masked-array vectorization numpy-ndarray numpy python
2 Answers

As far as I know, there is no option completely in numpy to do this. You could still use numpy to reduce the time for grabbing the indices, though.

data = [0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1]
n=3
def get_new_data(data,n):
  new_data = np.zeros(len(data))
  non_zero = np.argwhere(data).ravel()
  idx = non_zero[0]
  new_data[idx] =1
  idx += n
  for i in non_zero[1:]:
    if i > idx:
      new_data[i] = 1
      idx+=n
  return new_data
get_new_data(data, n)

A function like this should give you a better run time since you are not looping over the whole array.

If this is still not optimal to you, you can look at using numba, which works very well with numpy and is relatively easy to use.

user_41459410

You could do it like this:-

N = 3
data = [0, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1]
newData = data.copy()
i = 0
M = [0 for _ in range(N)]
while i < len(newData) - N:
    if newData[i] == 1:
        newData[i + 1:i + 1 + N] = M
        i += N
    i += 1

print(newData)
Admin1
Related