Fosa Fosa - 1 month ago 8
Python Question

Get # of contiguous hits and their first/last index in a NumPy array

Here is an itertools solution that returns a list of the lengths of each contiguous block. Here a contiguous block is a run of 1s without breaks in between. Is there a way to also have itertools return the index associated with each block?

import itertools
import numpy as np

stack = np.zeros((10,10))
stack[0] = 1
stack[5,:5] = 1
contiguous_hits = [ (sum( 1 for _ in group )) for row in stack for key, group in itertools.groupby(row) if key ]

Many thanks!

Answer Source

You could use np.where in conjunction with itertools:

In [477]: [[i, len(list(g))] for i, g in itertools.groupby(np.where(stack == 1)[0])]
Out[477]: [[0, 10], [5, 5]]

A list of lists is returned. In each list, the first item is the row index, the second item is the length of consecutive ones.