Skip to content
Advertisement

Replace consecutive identic elements in the beginning of an array with 0

I want to replace the N first identic consecutive numbers from an array with 0.

import numpy as np


x = np.array([1, 1, 1, 1, 2, 3, 1, 2, 3, 2, 2, 2, 3, 3, 3, 1, 1, 2, 2])

OUT -> np.array([0, 0, 0, 0 2, 3, 1, 2, 3, 2, 2, 2, 3, 3, 3, 1, 1, 2, 2])

Loop works, but what would be a faster-vectorized implementation?

i = 0
first = x[0]
while x[i] == first and i <= x.size - 1:
    x[i] = 0
    i += 1

Advertisement

Answer

You can use argmax on a boolean array to get the index of the first changing value.

Then slice and replace:

n = (x!=x[0]).argmax()  # 4
x[:n] = 0

output:

array([0, 0, 0, 0, 2, 3, 1, 2, 3, 2, 2, 2, 3, 3, 3, 1, 1, 2, 2])

intermediate array:

(x!=x[0])

#                            n=4
# [False False False False  True  True  True  True  True  True  True  True
#  True  True  True  True  True  True  True]
User contributions licensed under: CC BY-SA
3 People found this is helpful
Advertisement