python - pandas:仅当满足条件时才进行分组

Question

我有以下多索引 DataFrame:

        m     dist

a  a    2     5
   b    3     8
   c    4     12
   d    2     3
b  a    2     5
   b    3     8
   c    4     14
   d    2     27

我想根据算法计算一个新列 s。例如，对于 (a,a)，算法为:

选择具有相同 0 级索引的所有其他行: --> (a,a), (a,b), (a,c), (a,d)

从此选择中:选择 dist <= 自己的 dist 的行 --> (a,a) 和 (a,d)

获取该选择的第 m 列的总和 --> 2 + 2 = 4

这将产生以下数据框:

        m     dist  s

a  a    2     5     4
   b    3     8     7
   c    4     12    11
   d    2     3     2
b  a    2     5     2
   b    3     8     5
   c    4     14    9
   d    2     27    11

通过循环遍历行、为每一行创建一个新的数据框、选择具有正确距离的行并进行分组，这将很容易。但对于我当前的数据集来说，这太慢了。这可以通过一些不错的 groupby/lambda 魔法来完成吗？

Answer 1

您可以按 dist 对数据框进行排序，然后对列 m 执行 cumsum:

df['s'] = df.sort_values('dist').groupby(level=0).m.cumsum()

这假设 dist 列中没有重复项，如果有，则对按 dist 分组的列 s 进行进一步转换，并且第一级索引应满足:

df['s'] = df.groupby([df.index.get_level_values(0), 'dist']).s.transform('max')