numpy.arrayで連続してTrueであるものを探す

よくあるアレですがせっかく書いたんでコードをメモしておきます

やりたいこと

numpy.arrayについてN回連続してTrueのところだけ残します

やり方

1つずらしてandすることで2つ連続している起点を見つけ、1つずらしてorすることで終点までTrueで埋めます。

andの部分を2回やると3つ連続した起点、3回やると4つ連続した起点が見つけられますので、これを再起で書いてやれば任意の回数連続した部分を見つけられます。

1. 行列を準備

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. 連続しているところだけ残す

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)

これで以下のような出力が得られるはずです