Convert a list that intervals start and stop with [1,0,-1] into a step function [0, 1]-CodePudding

Say one is given a list that denotes when events occur. The start of an event is marked by a 1, and the end of the event is -1. Otherwise, for every time point sampled, it is 0. A short example might be:

event_list = [0,0,0,0,1,0,0,-1,0,0,0,0,0,0,0,0,1,0,0,0,0,0,-1,0,0,0,0,0,1,0,0,0,0,0,-1,0,0]

I would like to find a simple method to convert that list into a step function, where the start of an event is denoted by 1, the duration of the event is all 1's, and when the event ends it returns to 0. So for the same list as above, it would look something like:

step_list = [0,0,0,0,1,1,1,1,0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,0,0,0,0,0,1,1,1,1,1,1,1,0,0]

Now, one could brute force this by looping over the list and using if statements and a buffer to create a new list. But there has to be a far more elegant way than that. I was thinking something using the diff function in numpy, but that still doesn't quite work. Thanks for any suggestions!

CodePudding user response：

There may be a more elegant way (numpy can be surprising), but you can do this in two steps with .cumsum() and an indexed replacement:

import numpy as np
x = np.array(event_list)

result = x.cumsum()
result[x == -1] = 1

Output:

>>> np.array_equal(result, step_list)
True

This works by first taking the cumulative sum of the elements. The first 1 encountered sets the cumsum to 1. It's followed by zeros (which means the cumsum stays at 1) until a -1 is encountered, whereupon the cumsum is "reset" back to 0 (until the next 1 is encountered). This essentially "fills forward" the 1s up to, but not including the -1s, so the final step is to simply assign 1 to the indices where the original array has a value of -1.