python - 检索当前行和上一个/下一个正值之间的行数

Question

我有以下数据框:

   feature
0        1
1        0
2        0
3        0
4        0
5        1
6        0
7        1

我想创建一个 2 列，其中包含当前行与上一个和下一个正值之间的行数。输出数据框应该是这样的:

   feature    previous_feat        next_feat
0        1               NA                5
1        0                1                4
2        0                2                3
3        0                3                2
4        0                4                1
5        1                5                2
6        0                1                1
7        1                2               NA

我已经尝试过结合使用shift 和mask 方法，但我没能成功。请注意，它可能是行数或索引差异，这对我来说并不重要。 NA 值也一样，它可以是 NA 或 0。

import pandas as pd

df = pd.DataFrame({"feature": [1, 0, 0, 0, 0, 1, 0, 1]})

# df["previous_feat"] = df.shift().mask(df["feature"] != 0)

Answer 1

您可以使用 groupby.cumcount和 bool 掩码:

# form groups
g1 = df.loc[::-1, 'feature'].eq(1).cumsum()
g2 = df['feature'].eq(1).cumsum()

# mask first/last
m1 = g2.eq(1) & df['feature'].eq(1)
m2 = g1.eq(1) & df['feature'].eq(1)

# compute cumcount
df['previous_feat'] = df.groupby(g1).cumcount().add(1).mask(m1)
df['next_feat'] = df[::-1].groupby(g2).cumcount().add(1).mask(m2)

输出:

   feature  previous_feat  next_feat
0        1            NaN        5.0
1        0            1.0        4.0
2        0            2.0        3.0
3        0            3.0        2.0
4        0            4.0        1.0
5        1            5.0        2.0
6        0            1.0        1.0
7        1            2.0        NaN