Python pandas - 如何对封闭元素进行分组

Question

我有一个数据框，我需要在其中对距离不超过 1 的元素进行分组。例如，如果这是我的 df:

     group_number  val
0              1    5
1              1    8
2              1   12
3              1   13
4              1   22
5              1   26
6              1   31
7              2    7
8              2   16
9              2   17
10             2   19
11             2   29
12             2   33
13             2   62

所以我需要按 group_number 和 val 进行分组，其中 val 的值小于或等于 1。

所以，在这个例子中，行 2 和 3 将组合在一起，行 8 和 9 会聚在一起。

我尝试使用 diff 或相关函数，但我没有弄明白。

任何帮助将不胜感激!

Answer 1

使用 diff 是正确的方法 - 只需将它与 gt 和 cumsum 结合使用，您就拥有了自己的组。

想法是对大于阈值的差异使用累积和。大于阈值的差异将变为 True。相反，等于或低于您的阈值的差异将变为 False。对 bool 值的累积求和将使等于或低于您的阈值的差异保持不变，因此它们获得相同的组号。

max_distance = 1

df["group_diff"] = df.sort_values("val")\
                     .groupby("group_number")["val"]\
                     .diff()\
                     .gt(max_distance)\
                     .cumsum()

print(df)

    group_number    val group_diff
0   1               5   0
1   1               8   1
2   1               12  2
3   1               13  2
4   1               22  5
5   1               26  6
6   1               31  8
7   2               7   0
8   2               16  3
9   2               17  3
10  2               19  4
11  2               29  7
12  2               33  9
13  2               62  10

您现在可以在 group_number 和 group_diff 上使用 groupby 并查看结果组，结果如下:

grouped = df.groupby(["group_number", "group_diff"])
print(grouped.groups)

{(1, 0): Int64Index([0], dtype='int64'),
 (1, 1): Int64Index([1], dtype='int64'),
 (1, 2): Int64Index([2, 3], dtype='int64'),
 (1, 5): Int64Index([4], dtype='int64'),
 (1, 6): Int64Index([5], dtype='int64'),
 (1, 8): Int64Index([6], dtype='int64'),
 (2, 0): Int64Index([7], dtype='int64'),
 (2, 3): Int64Index([8, 9], dtype='int64'),
 (2, 4): Int64Index([10], dtype='int64'),
 (2, 7): Int64Index([11], dtype='int64'),
 (2, 9): Int64Index([12], dtype='int64'),
 (2, 10): Int64Index([13], dtype='int64')}

感谢@jezrael 提示避免使用新列以提高性能:

group_diff = df.sort_values("val")\
               .groupby("group_number")["val"]\
               .diff()\
               .gt(max_distance)\
               .cumsum()

grouped = df.groupby(["group_number", group_diff])