Search code examples
algorithmnumpy

Is there a hold value until condition function in numpy?


For example something like:

pre_hold_list = [-2,0,0,-1,0,0,0,3,0,0]
hold_condition = lambda x:x != 0
output = np.hold(pre_hold_list, hold_condition)
[-2,-2,-2,-1,-1,-1,-1,3,3,3] #result of output

Here the condition is that the current value is not zero the function will hold the value that this condition is met until the next value that meets this condition (i.e. it will hold -2 then -1 then 3).

Searching for np.hold() or np.step() does not give me anything on google.


Solution

  • You could use a trick here by using cumsum(…) [numpy-doc], and diff() [numpy-doc]:

    import numpy as np
    
    def hold(iterable, condition):
        cond = np.array(condition)
        vals = np.array(iterable)
        a = vals * cond
        a[cond] = np.diff(np.hstack(((0,), a[cond])))
        return a.cumsum()
    

    The first parameter is an iterable that contains the elements, the second parameter condition is an iterable of the same length with booleans.

    For example:

    >>> a
    array([-2,  0,  0, -1,  0,  0,  0,  3,  0,  0])
    >>> hold(a, a != 0)
    array([-2, -2, -2, -1, -1, -1, -1,  3,  3,  3])
    >>> hold(a, a != 0)
    array([-2, -2, -2, -1, -1, -1, -1,  3,  3,  3])
    

    The function works as follows. Furst we make a copy of the two iterables (and convert these to numpy arrays, if that is not yet the case). You can omit that if these are numpy arrays.

    Next we perform an elementwise multiplication, to make the values where the condition does not hold zero.

    Next we calculate the difference between each item where the condition holds and the next one, and we set that to a. Finally we can use the cummulative sum of a, since the .diff() ensured that this will result in the correct repetitions.