Python pandas - 在多列中的第一个非零值之前 NaN 全零的正确方法是什么？

Question

我有一个df列date ， a ， b和id 。 id是分组的， date值在转到新的id时重复。 在a和b列中，我想在每个id中的第一个非零值之前用nan替换 0 。 因此，使用以下数据：

df = pd.DataFrame({
    'date': ['2019-01-01', '2019-02-01', '2019-03-01', '2019-04-01', '2019-05-01']*3,
    'id': [0,0,0,0,0,1,1,1,1,1,2,2,2,2,2],
    'a': [0,0,10,40,20,0,0,0,50,90,0,0,0,0,0],
    'b': [0,0,0,123,345,0,0,555,0,666,0,0,0,0,30]
})

          date  id   a    b
0   2019-01-01   0   0    0
1   2019-02-01   0   0    0
2   2019-03-01   0  10    0
3   2019-04-01   0  40  123
4   2019-05-01   0  20  345
5   2019-01-01   1   0    0
6   2019-02-01   1   0    0
7   2019-03-01   1   0  555
8   2019-04-01   1  50    0
9   2019-05-01   1  90  666
10  2019-01-01   2   0    0
11  2019-02-01   2   0    0
12  2019-03-01   2   0    0
13  2019-04-01   2   0    0
14  2019-05-01   2   0   30

output 应该是

          date  id     a      b
0   2019-01-01   0   NaN    NaN
1   2019-02-01   0   NaN    NaN
2   2019-03-01   0  10.0    NaN
3   2019-04-01   0  40.0  123.0
4   2019-05-01   0  20.0  345.0
5   2019-01-01   1   NaN    NaN
6   2019-02-01   1   NaN    NaN
7   2019-03-01   1   NaN  555.0
8   2019-04-01   1  50.0    0.0
9   2019-05-01   1  90.0  666.0
10  2019-01-01   2   0.0    NaN
11  2019-02-01   2   0.0    NaN
12  2019-03-01   2   0.0    NaN
13  2019-04-01   2   0.0    NaN
14  2019-05-01   2   0.0   30.0

请注意，如果列中给定id的所有值都为零，则保留零。

我目前的解决方案是 2 个for -loops：一个用于列，一个用于id上的groupby对象； 我认为有改进空间的解决方案。 任何提示/帮助将不胜感激。

for col in ['a', 'b']:
    for i, grp in df.groupby('id'):
        min_idx = grp.index.min()
        non_z_idx = grp[grp[col] > 0].index.min()

        if not np.isnan(non_z_idx):
            df.loc[min_idx:non_z_idx - 1, col] = np.nan

Answer 1

使用 2 个面具和df.where

m1 = df[['a','b']].ne(0).groupby(df.id).cummax()
m2 = df[['a','b']].eq(0).groupby(df.id).transform('all')

df[['a','b']] = df[['a','b']].where(m1 | m2)

Out[88]:
          date  id     a      b
0   2019-01-01   0   NaN    NaN
1   2019-02-01   0   NaN    NaN
2   2019-03-01   0  10.0    NaN
3   2019-04-01   0  40.0  123.0
4   2019-05-01   0  20.0  345.0
5   2019-01-01   1   NaN    NaN
6   2019-02-01   1   NaN    NaN
7   2019-03-01   1   NaN  555.0
8   2019-04-01   1  50.0    0.0
9   2019-05-01   1  90.0  666.0
10  2019-01-01   2   0.0    NaN
11  2019-02-01   2   0.0    NaN
12  2019-03-01   2   0.0    NaN
13  2019-04-01   2   0.0    NaN
14  2019-05-01   2   0.0   30.0

如果您不想要 2 个 groupby，您可以将一个 groupby 与apply

m = df[['a','b']].ne(0).groupby(df.id).apply(lambda x: x.cummax() | ~x.any())
df[['a','b']] = df[['a','b']].where(m)