首页 > 解决方案 > Pandas:我可以过滤数据框以仅获取彼此之间相差 50% 的行吗?

问题描述

以下是一个示例数据框。我的实际数据集有 30k 行。

df = pd.DataFrame({'Account': [30, 30, 30, 30, 30, 30, 30, 40, 40, 40],  
                   'Start': [2, 2, 2, 2, 2, 3, 3, 1, 1, 1],  
                   'Amount' : [500, 600, 800, 200, 700, 10, 800, 10, 50, 70]})

   Account Start  Amount
0       30     2     500
1       30     2     600
2       30     2     800
3       30     2     200
4       30     2     700
5       30     3      10
6       30     3     800
7       40     1      10
8       40     1      50
9       40     1      70

我想查找第 1 行中的金额与第 2 行中的金额相差 ± 50% 的所有行(按帐户和开始分组)。我期待结果看起来像这样。

   Account Start  Amount
0       30     2     500
1                    600
2                    800
8       40     1      50
9                     70

第 3 行被排除,因为第 3 行中的 200 小于第 2 行中的金额以及第 3
行中的金额的 50%。第 4 行被排除,因为它是 start = 2 中的最后一个元素,并且前一行也被排除在外.
同样,第 5 行和第 6 行也被排除在外。
第 7 行被排除在外,因为 10 少于第 8 行中金额的 50%。

PS:在最终的数据集中,每组Account和Start至少应该有4行。
有没有办法有效地做到这一点?

标签: pythonpandaspandas-groupby

解决方案


我们使用pct_change,检查它是否在 -50% 和 50% 之间。因为你想要成对的行,我们需要检查这个掩码或移位掩码(在我们计算 pct_change 的相反方向上移动)。我们将把这个函数分别应用于每个组。

def keep_within_pct(gp, shift=1, pcts=(-0.5, 0.5)):
    m = gp['Amount'].pct_change(-shift).between(*pcts)
    return gp[m | m.shift(shift).fillna(False)]

df.groupby(['Account', 'Start'], group_keys=False).apply(keep_within_pct)

   Account  Start  Amount
0       30      2     500
1       30      2     600
2       30      2     800
8       40      1      50
9       40      1      70

推荐阅读