python - Fill zero values of 1d numpy array with last non-zero values

Question

Let's say we have a 1d numpy array filled with some int values. And let's say that some of them are 0.

Is there any way, using numpy array's power, to fill all the 0 values with the last non-zero values found?

for example:

arr = np.array([1, 0, 0, 2, 0, 4, 6, 8, 0, 0, 0, 0, 2])
fill_zeros_with_last(arr)
print arr

[1 1 1 2 2 4 6 8 8 8 8 8 2]

A way to do it would be with this function:

def fill_zeros_with_last(arr):
    last_val = None # I don't really care about the initial value
    for i in range(arr.size):
        if arr[i]:
            last_val = arr[i]
        elif last_val is not None:
            arr[i] = last_val

However, this is using a raw python for loop instead of taking advantage of the numpy and scipy power.

If we knew that a reasonably small number of consecutive zeros are possible, we could use something based on numpy.roll. The problem is that the number of consecutive zeros is potentially large...

Any ideas? or should we go straight to Cython?

Disclaimer:

I would say long ago I found a question in stackoverflow asking something like this or very similar. I wasn't able to find it. :-(

Maybe I missed the right search terms, sorry for the duplicate then. Maybe it was just my imagination...

Answer 1

Here's a solution using np.maximum.accumulate:

def fill_zeros_with_last(arr):
    prev = np.arange(len(arr))
    prev[arr == 0] = 0
    prev = np.maximum.accumulate(prev)
    return arr[prev]

We construct an array prev which has the same length as arr, and such that prev[i] is the index of the last non-zero entry before the i-th entry of arr. For example, if:

>>> arr = np.array([1, 0, 0, 2, 0, 4, 6, 8, 0, 0, 0, 0, 2])

Then prev looks like:

array([ 0,  0,  0,  3,  3,  5,  6,  7,  7,  7,  7,  7, 12])

Then we just index into arr with prev and we obtain our result. A test:

>>> arr = np.array([1, 0, 0, 2, 0, 4, 6, 8, 0, 0, 0, 0, 2])
>>> fill_zeros_with_last(arr)
array([1, 1, 1, 2, 2, 4, 6, 8, 8, 8, 8, 8, 2])

Note: Be careful to understand what this does when the first entry of your array is zero:

>>> fill_zeros_with_last(np.array([0,0,1,0,0]))
array([0, 0, 1, 1, 1])