Python 3: test for contiguous numbers in one line

Question

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?

score 0 · Accepted Answer · answered Oct 08 '21 at 23:50

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]]

Python 3: test for contiguous numbers in one line

1 Answers1