Pandas中apply能否跨组调用数据？求组间元素存在性判断最优方案

需求说明 #

对DataFrame中每个组内的元素，判断其是否存在于下一个组（按组在数据中的出现顺序，非数值排序），最后一组的所有元素返回False。

示例输入 #

import pandas as pd

df = pd.DataFrame({'group': [ 0,   1,   1,   0,   2 ], 
                     'val': ['a', 'b', 'a', 'c', 'c']})
grouped = df.groupby('group')

预期输出 #

print(result)
0     True
1    False
2    False
3    False
4    False
Name: val, dtype: bool

原实现（可优化） #

keys = list(grouped.groups.keys())

iterator_keys = iter(keys[1:])
def f(ser):
    if ser.name == keys[-1]:
        return ser.isin([])
    next_key = next(iterator_keys)
    return ser.isin(grouped.get_group(next_key)['val'])
result = grouped['val'].apply(f)

优化实现 #

原代码依赖外部迭代器，存在状态依赖（多次执行会出错），且逻辑不够直观。下面是更简洁可靠的实现：

import pandas as pd

df = pd.DataFrame({'group': [0, 1, 1, 0, 2], 'val': ['a', 'b', 'a', 'c', 'c']})

# 获取组的实际出现顺序（按数据中首次出现的顺序）
group_order = df['group'].unique().tolist()

# 构建每个组对应的下一组值集合
next_group_vals = {}
for idx, group in enumerate(group_order):
    if idx == len(group_order) - 1:
        # 最后一组对应空集合
        next_group_vals[group] = set()
    else:
        # 取当前组的下一组的所有val值
        next_group = group_order[idx + 1]
        next_group_vals[group] = set(df[df['group'] == next_group]['val'])

# 用transform实现向量化判断，避免状态依赖
result = df.groupby('group')['val'].transform(lambda ser: ser.isin(next_group_vals[ser.name]))

print(result)

优化点说明 #

1. 无状态逻辑：避免了原代码中迭代器的状态依赖，多次执行也不会出问题
2. 提前预计算：一次性构建好所有组对应的下一组值集合，减少重复计算，提升效率
3. 更符合pandas风格：使用transform替代apply，逻辑更直观，可读性更强

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