如何基于子集关系高效分组含集合列的Pandas DataFrame？

针对你提出的需求，这里提供一套优化方案，能大幅降低实际运行时间，尤其是数据量较大时：

核心优化思路 #

大集合不可能是小集合的子集，因此我们可以先处理大集合，用“代表集合”映射分组，避免两两比较的O(n²)复杂度；同时用位运算替代集合的子集判断，进一步加速操作。

完整实现代码 #

import pandas as pd

# 生成示例DataFrame
df = pd.DataFrame({
    'A': [{'A', 'B'}, {'A', 'B', 'C', 'E'}, {'B', 'D'}, {'C', 'B'}, {'A', 'B', 'D'}, {'X'}],
    'B': [111, 222, 333, 444, 555, 666]
})

# 步骤1：将集合转为可哈希的frozenset，并计算集合长度
df['frozen_A'] = df['A'].apply(frozenset)
df['len_A'] = df['A'].apply(len)

# 步骤2：将所有元素映射为bit位，把集合转为整数mask（加速子集判断）
all_elements = set().union(*df['A'])
element_to_bit = {elem: 1 << idx for idx, elem in enumerate(all_elements)}
df['bitmask'] = df['frozen_A'].apply(lambda s: sum(element_to_bit[elem] for elem in s))

# 步骤3：按集合长度降序排序，优先处理大集合
sorted_df = df.sort_values(by='len_A', ascending=False)

# 步骤4：维护代表集合与组号的映射，给每个原索引分配组号
group_map = {}  # key: 代表集合的bitmask, value: 组号
current_group = 0
index_to_group = {}

for original_idx, row in sorted_df.iterrows():
    current_mask = row['bitmask']
    assigned = False
    # 仅与已有的代表集合比较，无需遍历所有集合
    for rep_mask, group_num in group_map.items():
        # 位运算判断子集：小集合mask & 大集合mask == 小集合mask
        if (current_mask & rep_mask) == current_mask:
            index_to_group[original_idx] = group_num
            assigned = True
            break
    # 如果不是任何代表集合的子集，自身成为新的代表集合
    if not assigned:
        group_map[current_mask] = current_group
        index_to_group[original_idx] = current_group
        current_group += 1

# 步骤5：将组号映射回原DataFrame
df['group'] = df.index.map(index_to_group)

# 查看分组结果
print("分组结果：")
for group_num in sorted(df['group'].unique()):
    indices = df[df['group'] == group_num].index.tolist()
    print(f"第{group_num+1}组：索引{indices}")

# 按组聚合操作
grouped = df.groupby('group')

运行结果 #

分组结果：
第1组：索引[0, 1, 3]
第2组：索引[2, 4]
第3组：索引[5]

效率说明 #

1. 时间复杂度：排序步骤为O(n log n)，遍历过程中每个集合仅与代表集合比较（而非所有n个集合），代表集合的数量通常远小于n（示例中仅3个）。位运算判断子集是O(1)操作，比原生集合issubset()快数倍。
2. 适用场景：当数据量较大、存在较多子集关系时，效率提升尤为明显；即使所有集合互不包含，最坏情况也只是O(n²)，但实际运行速度仍因位运算优化快于原生集合比较。

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