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

Question

我有一个 df 带列 date, a, b 和一个 id. 该 id 是分组和 date 当进入一个新版本时，数值会重复。id. 在专栏中 a 和 b 我想把0换成 nan 之前 的第一个非零值，在每个 id. 所以有了下面的数据。

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

输出应该是

          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-循环：一个用于柱子，一个用于循环。groupby 物体上 id;一个解决方案，我相信有改进的空间。任何提示和帮助将是非常感激的。

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

如果您不需要两个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)