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