高效实现Pandas DataFrame按行集合差的因式分解方法问询

核心思路 #

利用Pandas原生方法结合集合操作，简化逻辑同时提升效率，避免冗余循环或复杂处理。

1. 准备示例数据 #

import pandas as pd

# 全局地点列表
list_locations = ["A", "B", "C", "D", "E"]
# 聚合后的输入DataFrame
df = pd.DataFrame({
    "product_group": ["PG1", "PG2", "PG3"],
    "location_list": [["A", "B"], ["C"], ["A", "D", "E"]]
})

2. 计算rest_of_world_location_list #

先将全局地点转为集合（提升查找效率），通过apply结合集合操作完成集合差计算，可按需选择是否保留原地点顺序：

方式1：不保留顺序（速度更快）

直接利用集合差操作，时间复杂度更低：

global_loc_set = set(list_locations)
df["rest_of_world_location_list"] = df["location_list"].apply(
    lambda locs: list(global_loc_set - set(locs))
)

方式2：保留list_locations的原始顺序

通过列表推导筛选不在当前分组地点中的元素：

global_loc_set = set(list_locations)
df["rest_of_world_location_list"] = df["location_list"].apply(
    lambda locs: [loc for loc in list_locations if loc not in set(locs)]
)

3. 因式分解生成rest_of_world_index #

由于列表是不可哈希对象，先转为元组后使用pd.factorize生成唯一索引：

# 转换为可哈希的元组后因式分解
df["rest_of_world_index"] = pd.factorize(df["rest_of_world_location_list"].apply(tuple))[0]

最终输出 #

运行上述代码后，DataFrame结果如下：

product_group location_list rest_of_world_location_list  rest_of_world_index
0           PG1       [A, B]                [C, D, E]                    0
1           PG2          [C]                [A, B, D, E]                  1
2           PG3    [A, D, E]                [B, C]                      2

效率亮点 #

• 集合操作将元素查找的时间复杂度从O(n)降至O(1)，大幅提升大数据集处理速度；
• pd.factorize是Pandas原生高效方法，比手动映射唯一值的循环逻辑快数倍；
• 整体逻辑简洁，仅需3步即可完成转换，避免冗余的中间变量。

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