How to build a mask (True or False) for consecutive zeroes in a row in a pandas dataframe?

user_9784713 0
Viewed 183

I have a dataframe and want to create a mask, marking the occurrence of consecutive zeros (3 or more) in it.

I want to mark as True the occurrence of zeros when they happen three or more times consecutively. That is, three or more zeros in a row, without any type of interruption (such as other numbers or NaN). Therefore, the array [0,0,0,1,2,3,4] becomes [True, True, True, False, False, False, False] and the array [0,1,0,2,0,0,0] becomes [False,False,False,False,True,True,True].

The example bellow illustrates better the input and output.

Input:

col1    col2    col3    col4    col5
0       0       0       0       0
0       0       5       0       5
0       0       0       3       3
0       0       0       NaN     0

I want the result to be something like

col1  col2  col3  col4  col5
True  True  True  True  True
False False False False False
True  True  True  False False
True  True  True  False False

Currently I am doing this iterating over the rows, but it is somewhat slow (this dataframe has actually over 1 million rows).

Is there any way I can accomplish this avoiding for loops (iterrows, apply, etc) over the rows?

Thanks!

pandas numpy python
1 Answers

I'll show a numpy solution and leave it to you to transform to pandas.

a = np.array([
    [0, 0, 0],
    [0, 0, 0],
    [0, 5, 0],
    [0, 0, 3],
    [0, 5, 3]]).T

First, mask the array and pad the mask with zeros on each row:

z = np.zeros((a.shape[0], 1), dtype=bool)
mask = np.concatenate((z, a == 0, z), axis=1)

By finding the location where the mask changes value, you can find the length of each run. The padding in the previous step ensures that the first change will be "on" and the last will be "off":

locs = np.nonzero(np.diff(mask, axis=1))

Now is the slightly tricky part. Normally you would do this in 1D, but here is 2D. However, the padding ensures that your run lengths are still accurate:

run_lengths = locs[1][1::2] - locs[1][::2]

Now you want to turn off the elements of mask that represent runs of fewer than n=3 elements. With the information you have, it may be easier to just create a new mask directly.

First mask out the runs that are too short:

valid_runs = np.flatnonzero(run_lengths >= 3)

Then make the output as a int8, which we will later view as a bool:

result = np.zeros(a.shape, dtype=np.int8)

Set the first element of each valid run to 1, and one past the end to -1:

v = 2 * valid_runs
result[locs[0][v], locs[1][v]] = 1
v += 1
v = v[locs[1][v] < result.shape[1]]
result[locs[0][v], locs[1][v]] = -1

You can take the cumulative sum of the result in-place, and view it as a boolean mask (since int8 and bool_ have the same size):

result = np.cumsum(result, axis=1, out=result).view(bool)

The final result is:

array([[ True,  True,  True,  True,  True],
       [False, False, False, False, False],
       [ True,  True,  True, False, False]])

TL;DR

Here is a fully generalized solution that will work across any axis in a multi-dimensional array, with an arbitrary selection for minimum run length (n) and number (k):

def mask_consecutive(a, k=0, n=3, axis=-1):
    shape = list(a.shape)
    shape[axis] = 1;
    z = np.zeros(shape, dtype=bool)
    mask = np.concatenate((z, a == k, z), axis=axis)
    locs = np.argwhere(np.diff(mask, axis=axis))
    run_lengths = locs[1::2, axis] - locs[::2, axis]
    valid_runs = np.flatnonzero(run_lengths >= n)
    result = np.zeros(a.shape, dtype=np.int8)
    v = 2 * valid_runs
    result[tuple(locs[v, :].T)] = 1
    v += 1
    v = v[locs[v, axis] < result.shape[axis]]
    result[tuple(locs[v, :].T)] = -1
    return np.cumsum(result, axis=axis, out=result).view(bool)
user_29887300
Related