Count number of zeros after every non-zero clusters in python

Viewed 114

I want to count the number of zeros after every batch of non-zero numbers in a dataframe. For example, see this dataframe. The number of zeros in column A after the first non-zero cluster is 2. Then after the next cluster, the number of zeros is 4 and after the final cluster the number of zeros is 2. So the output should be [2,4,2].

Time A  B  C
60   5  4  3
120  4  5  2
180  5  5  5
240  0  0  0
300  0  0  0 
360  6  7  8
420  5  6  7
480  0  0  0
540  0  0  0 
600  0  0  0
660  0  0  0
720  5  6  6
780  5  6  6
840  0  0  0
900  0  0  0 
2 Answers

Using itertools.groupby:

>>> from itertools import groupby
>>> [len(list(group)) for key, group in groupby(df.A) if key == 0]
[2, 4, 2]

The elements themselves are used in grouping which means each distinct element will start a new group; so whenever a 0 starts a group as the key, we take the length of the (lazy) group as the result.

for c in df.loc[:, "A":]:
    m = df[c] == 0
    print(
        c,
        df[c]
        .groupby((m != m.shift(1)).cumsum())
        .apply(lambda x: len(x) if x.iat[0] == 0 else np.nan)
        .dropna()
        .tolist(),
    )

Prints:

A [2.0, 4.0, 2.0]
B [2.0, 4.0, 2.0]
C [2.0, 4.0, 2.0]
Related