如何基于绝对差过滤Pandas DataFrame并拆分结果？

原始数据 #

创建DataFrame的代码：

import pandas as pd
d1 = {'id': ["car", "car", "car", "plane", "plane", "car"], 'value': [1, 1.2, 5, 6, 1.3, 0.8]}
df1 = pd.DataFrame(data=d1)

原始数据内容：

id      value
0   car     1.0
1   car     1.2
2   car     5.0
3   plane   6.0
4   plane   1.3
5   car     0.8

需求 #

将原始DataFrame拆分为两个结果集：

• 结果集1：包含所有属于某一子集的行，该子集内同一ID的所有value两两绝对差均小于1
• 结果集2：包含剩余的所有行

预期结果：

结果集1（df2） #

d2 = {'id': ["car", "car", "car"], 'value': [1, 1.2, 0.8]}
df2 = pd.DataFrame(data=d2)

id      value
0   car     1.0
1   car     1.2
5   car     0.8

结果集2（df3） #

d3 = {'id': ["car", "plane", "plane"], 'value': [5, 6, 1.3]}
df3 = pd.DataFrame(data=d3)

2   car     5.0
3   plane   6.0
4   plane   1.3

尝试的代码（未达预期） #

unique_list = []
def unique_2(df):
    for id_1, value_1 in zip(df["id"], df["value"]):
        for id_2, value_2 in zip(df["id"], df["value"]):
            if id_1 == id_2:
                if abs(value_1-value_2) > 0.01: 
                    x = True
                    unique_list.append(x)
                else:
                    x = False
                    unique_list.append(x)
            else:
                pass

正确实现方案 #

思路 #

按ID分组后，对每组的value排序，通过滑动窗口找到最长的连续区间（区间内最大值与最小值的差小于1），该区间内的行即为符合条件的行；组内其余行归入剩余结果集。

代码实现 #

import pandas as pd

def filter_group(group):
    # 对组内value排序，保留原始索引
    sorted_vals = group['value'].sort_values().reset_index()
    n = len(sorted_vals)
    max_len = 0
    best_start = 0
    
    # 滑动窗口寻找最长符合条件的区间
    for i in range(n):
        for j in range(i, n):
            if sorted_vals['value'].iloc[j] - sorted_vals['value'].iloc[i] < 1:
                current_len = j - i + 1
                if current_len > max_len:
                    max_len = current_len
                    best_start = i
            else:
                break  # 排序后后续值更大，差值只会递增
    
    # 分割符合条件与不符合条件的行
    if max_len >= 1:
        valid_indices = sorted_vals['index'].iloc[best_start:best_start+max_len].tolist()
        valid = group.loc[valid_indices]
        invalid = group.drop(valid_indices)
        return valid, invalid
    else:
        return pd.DataFrame(), group

# 原始数据
d1 = {'id': ["car", "car", "car", "plane", "plane", "car"], 'value': [1, 1.2, 5, 6, 1.3, 0.8]}
df1 = pd.DataFrame(data=d1)

# 分组处理并合并结果
valid_dfs = []
invalid_dfs = []
for _, group in df1.groupby('id'):
    valid, invalid = filter_group(group)
    valid_dfs.append(valid)
    invalid_dfs.append(invalid)

df2 = pd.concat(valid_dfs).sort_index()
df3 = pd.concat(invalid_dfs).sort_index()

# 输出结果
print("结果集1：")
print(df2)
print("\n结果集2：")
print(df3)

输出结果 #

结果集1：
    id  value
0  car    1.0
1  car    1.2
5  car    0.8

结果集2：
      id  value
2    car    5.0
3  plane    6.0
4  plane    1.3

内容的提问来源于stack exchange，提问作者Limmi