Split a Pandas dataframe into multiple dataframes based on the value of a column-CodePudding

How can I split a Pandas dataframe into multiple dataframes based on the value in a column?

df = pd.DataFrame({'A':[4,5,0,0,5,0,0,4],
                   'B':[7,8,0,0,4,0,0,0],
                   'C':[1,3,0,0,7,0,0,0]}, columns = ['A','B','C'])

df["sum"] = df.sum(axis=1)
df["Rolling_sum"] = df["sum"].rolling(2, min_periods=1).sum()

The resulting dataframe is:

    A   B   C   sum     Rolling_sum
0   4   7   1   12  12.0
1   5   8   3   16  28.0
2   0   0   0   0   16.0
3   0   0   0   0   0.0
4   5   4   7   16  16.0
5   0   0   0   0   16.0
6   0   0   0   0   0.0
7   4   0   0   4   4.0

I want to split the dataframe into multiple dataframe based on the occurrence of 0 in the Rolling_sum column.

Expected result:

Dataframe 1:

    A   B   C   sum     Rolling_sum
0   4   7   1   12  12.0
1   5   8   3   16  28.0
2   0   0   0   0   16.0

Dataframe 2:

    A   B   C   sum     Rolling_sum
4   5   4   7   16  16.0
5   0   0   0   0   16.0

Dataframe 3:

    A   B   C   sum     Rolling_sum
7   4   0   0   4   4.0

I'm not sure what condition(s) I can use to split the dataframe.

CodePudding user response：

You can do cumsum create the groupby key then groupby

d = {x : y for x , y in df.loc[df['Rolling_sum'].ne(0)].groupby(df['Rolling_sum'].eq(0).cumsum())}
d
Out[260]: 
{0:    A  B  C  sum  Rolling_sum
 0  4  7  1   12         12.0
 1  5  8  3   16         28.0
 2  0  0  0    0         16.0, 1:    A  B  C  sum  Rolling_sum
 4  5  4  7   16         16.0
 5  0  0  0    0         16.0, 2:    A  B  C  sum  Rolling_sum
 7  4  0  0    4          4.0}

CodePudding user response：

If you want to keep the zeros, np.split is easiest:

np.split(df, df.index[df['Rolling_sum'] == 0])

# [   A  B  C  sum  Rolling_sum
#  0  4  7  1   12         12.0
#  1  5  8  3   16         28.0
#  2  0  0  0    0         16.0,
#
#     A  B  C  sum  Rolling_sum
#  3  0  0  0    0          0.0
#  4  5  4  7   16         16.0
#  5  0  0  0    0         16.0,
#
#     A  B  C  sum  Rolling_sum
#  6  0  0  0    0          0.0
#  7  4  0  0    4          4.0]

If you want to ignore the zeros, adjust the cut points to account for the missing rows:

cuts = df.index[df['Rolling_sum'] == 0] # [3, 6]
cuts -= np.arange(len(cuts))   1        # [2, 4]
np.split(df[df['Rolling_sum'] != 0], cuts)

# [   A  B  C  sum  Rolling_sum
#  0  4  7  1   12         12.0
#  1  5  8  3   16         28.0,
#
#     A  B  C  sum  Rolling_sum
#  2  0  0  0    0         16.0
#  4  5  4  7   16         16.0,
#
#     A  B  C  sum  Rolling_sum
#  5  0  0  0    0         16.0
#  7  4  0  0    4          4.0]

Note that df.index[df['Rolling_sum'] == 0] assumes you have a default range index
If you have a custom index, use np.where(df['Rolling_sum'] == 0)[0] instead