Search code examples
pythonnumpypython-itertools

How to group blocks of identical Booleans?


Suppose I have the following list:

a = [True, True, True, False, False, False, False, True, True]

How could get them best grouped by either returning only 0,3, 7 or a grouping like the following?

[True, True, True]
[False, False, False, False]
[True, True]

Background: I am trying to find plateaus in my NumPy arrays, and while setting the derivative to zero is a good start, I still need to sort the array into chunks. I think this basically boils down to the problem above.

I looked up NumPy and itertools (trying to get a solution from the question NumPy grouping using itertools.groupby performance) but I did not succeed. I guess one might use a combination of itertools.takewhile and filtfalse (see the documentation here), but I am out of my depth there. Or maybe I am just thinking way too complicated.


Solution

  • We could get the indices with a sliced array comparison, and it should be good with performance with large size lists/arrays -

    a_ext = np.r_[~a[0],a]
    out = np.flatnonzero(a_ext[:-1]!=a_ext[1:])
    

    As a one-liner, we could use np.diff + np.flatnonzero -

    np.flatnonzero(np.diff(np.r_[~a[0],a]))
    # compact alternative : np.where(np.diff(np.r_[~a[0],a]))[0]