Python 3: test for contiguous numbers in one line-CodePudding

I've got an algorithm that works fine, but I was hoping to implement it another way for personal satisfaction.

In brief: I have some array obj_level that is a boolean mask indicating a coordinate where an object is present, so something like

 obj_level = [ 0, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0 , 1, 1 ]

I want to identify the base and the top of the object(s).

I first do that by producing an array containing the indices where the object is nonzero, obj_idx. If the array isn't empty, the first value is appended to base.

Then I loop through and test if the value of the index 1 is equal to the next object in obj_idx. If it is, no edge has been found, continue.

Otherwise, I've found an edge, so I append to base and top. I can make the inference that there is both a top and a base at a found edge because there are more values in obj_idx as I'm testing against obj_idx[i 1].

Finally, I append the last value to top, as the object must have a top if there is a base.

    base = []
    top = []
    obj_idx = np.flatnonzero(obj_level)
    if obj_idx.size > 0:
        base . append(obj_idx[0])
        for i,idx in enumerate(obj_idx[:-1]):
            if idx 1 == obj_idx[i 1]:
                continue
            else:
                top.append(idx)
                base.append(obj_idx[i 1])
        top.append(obj_idx[-1])

I'd like to do this in fewer lines. something like:

    base = [
           idx   1 == obj_idx[i 1] or idx 1
           for i,idx in enumerate(obj_idx[:-1])
           ]
    top =  [
           (idx 1 == obj_idx[i 1] or idx
           for i,idx in enumerate(obj_idx[:-1])
           ]
    np.insert(base,0,obj_idx[0])
    np.insert(top,-1,obj_idx[-1])

But I end up with a mixed array something like [True, True, 3, True, True]

Is there simpler way to do this than by extracting ints from the mixed array?

CodePudding user response：

You can just use the 'if' clause in list comprehension, as described here if/else in a list comprehension, to filter the source iterable. Something like (I avoided using numpy to concatenate the value just to make up a one-liner, but of course it wouldn't change anything):

base = [obj_idx[0]]   [
   obj_idx[i 1] for i in range(1, len(obj_idx)-1) if (obj_idx[i 1] != obj_idx[i]   1)
   ]
top = [
   obj_idx[i] for i in range(1, len(obj_idx)-1) if (obj_idx[i 1] != obj_idx[i]   1)
   ]   [obj_idx[-1]]