Search numpy.array for consecutive True

It's a common thing, but I wrote it so much so I'll write down the code

Thing you want to do

Leave only True for numpy.array N times in a row

manner

By shifting and and, you can find two consecutive starting points, and by shifting or shifting by one, you can fill up to the end point with True.

If you do the and part twice, you will find 3 consecutive starting points, and if you do it 3 times, you will find 4 consecutive starting points, so if you write this again, you can find the continuous part any number of times.

1. Prepare the matrix

python

import numpy as np
arr = np.array([[True,  False, False, True,  False, False, False],
                [False, False, True,  True,  True,  True,  False],
                [False, True,  True,  False, True,  False, True ],
                [False, False, False, False, True,  True,  False],
                [True,  True,  False, False, False, False, False]])
arr

2. Leave only continuous areas

python

def is_continuous(arr, left_cnt):
    if left_cnt >= 1:
        arr2 = arr[:, :-1] & arr[:, 1:]
        arr2 = is_continuous(arr2, left_cnt - 1)
    else:
        arr2 = arr
    return arr2

def check_continuous(arr, continuous_cnt):
    arr2 = is_continuous(arr, continuous_cnt - 1)
    arr3 = np.tile(np.full(continuous_cnt - 1, False), arr2.shape[0]).reshape(-1, arr2.shape[0]).T
    arr2 = np.hstack([arr2, arr3])
    for _ in range(continuous_cnt - 1):
        arr2[:, 1:] = arr2[:, 1:] | arr2[:, :-1]
    return arr2

print(arr)
check_continuous(arr, 2)

You should now get output similar to the following