Home > Back-end >  In a 2D Numpy Array find max streak of consecutive 1's
In a 2D Numpy Array find max streak of consecutive 1's

Time:02-02

I have a 2d numpy array like so. I want to find the maximum consecutive streak of 1's for every row.

a = np.array([[1, 1, 1, 1, 1],
              [1, 0, 1, 0, 1],
              [1, 1, 0, 1, 0],
              [0, 0, 0, 0, 0],
              [1, 1, 1, 0, 1],
              [1, 0, 0, 0, 0],
              [0, 1, 1, 0, 0],
              [1, 0, 1, 1, 0],
              ]
             )

Desired Output: [5, 1, 2, 0, 3, 1, 2, 2]

I have found the solution to above for a 1D array:

a = np.array([1, 1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0])
d = np.diff(np.concatenate(([0], a, [0])))
np.max(np.flatnonzero(d == -1) - np.flatnonzero(d == 1))
> 4

On similar lines, I wrote the following but it doesn't work.

d = np.diff(np.column_stack(([0] * a.shape[0], a, [0] * a.shape[0])))
np.max(np.flatnonzero(d == -1) - np.flatnonzero(d == 1))

CodePudding user response:

The 2D equivalent of you current code would be using pad, diff, where and maximum.reduceat:

# pad with a column of 0s on left/right
# and get the diff on axis=1
d = np.diff(np.pad(a, ((0,0), (1,1)), constant_values=0), axis=1)

# get row/col indices of -1
row, col = np.where(d==-1)

# get groups of rows
val, idx = np.unique(row, return_index=True)

# subtract col indices of -1/1 to get lengths
# use np.maximum.reduceat to get max length per group of rows
out = np.zeros(a.shape[0], dtype=int)
out[val] = np.maximum.reduceat(col-np.where(d==1)[1], idx)

Output: array([5, 1, 2, 0, 3, 1, 2, 2])

Intermediates:

np.pad(a, ((0,0), (1,1)), constant_values=0)

array([[0, 1, 1, 1, 1, 1, 0],
       [0, 1, 0, 1, 0, 1, 0],
       [0, 1, 1, 0, 1, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 1, 1, 1, 0, 1, 0],
       [0, 1, 0, 0, 0, 0, 0],
       [0, 0, 1, 1, 0, 0, 0],
       [0, 1, 0, 1, 1, 0, 0]])


np.diff(np.pad(a, ((0,0), (1,1)), constant_values=0), axis=1)

array([[ 1,  0,  0,  0,  0, -1],
       [ 1, -1,  1, -1,  1, -1],
       [ 1,  0, -1,  1, -1,  0],
       [ 0,  0,  0,  0,  0,  0],
       [ 1,  0,  0, -1,  1, -1],
       [ 1, -1,  0,  0,  0,  0],
       [ 0,  1,  0, -1,  0,  0],
       [ 1, -1,  1,  0, -1,  0]])


np.where(d==-1)

(array([0, 1, 1, 1, 2, 2, 4, 4, 5, 6, 7, 7]),
 array([5, 1, 3, 5, 2, 4, 3, 5, 1, 3, 1, 4]))


col-np.where(d==1)[1]

array([5, 1, 1, 1, 2, 1, 3, 1, 1, 2, 1, 2])


np.unique(row, return_index=True)

(array([0, 1, 2, 4, 5, 6, 7]),
 array([ 0,  1,  4,  6,  8,  9, 10]))


out = np.zeros(a.shape[0], dtype=int)

array([0, 0, 0, 0, 0, 0, 0, 0])


out[val] = np.maximum.reduceat(col-np.where(d==1)[1], idx)

array([5, 1, 2, 0, 3, 1, 2, 2])
  • Related