在数据框中找到列表交集而不是逐行查找的更快方法？答案

【问题标题】：Faster way of finding intersection of lists in a dataframe rather than going row by row?在数据框中找到列表交集而不是逐行查找的更快方法？
【发布时间】：2021-02-22 20:11:51
【问题描述】：

我有两个 pandas 数据框，它们都有相同的两列，其中每个元素都是一个列表。比如：

df1:

|---------------------|------------------|
|        col_a        |        col_b     |
|---------------------|------------------|
|   [A1, A2, A3, A4]  |        [B1]      |
|---------------------|------------------|
|   [A1, A2, A3, A4]  |      [B2, B3]    |
|---------------------|------------------|
|         [A5]        |      [B2, B3]    |
|---------------------|------------------|

df2:

|---------------------|------------------|
|        col_a        |        col_b     |
|---------------------|------------------|
|       [A2, A4]      |      [B1, B2]    |
|---------------------|------------------|
|         [A5]        |      [B2, B4]    |
|---------------------|------------------|

每一行代表一组组合，即df1的第1行代表A1-B1、A2-B1、A3-B1和A4-B1的组合。我想要做的是逐行获取一组独特的组合，这些组合同时存在于 df1 和 df2 中，但仍保留在此列表表示中。我不希望每一行都代表一个独特的组合。

示例的输出如下所示：

输出：

|---------------------|------------------|
|        col_a        |        col_b     |
|---------------------|------------------|
|        [A2, A4]     |        [B1]      |
|---------------------|------------------|
|        [A2, A4]     |        [B2]      |
|---------------------|------------------|
|         [A5]        |        [B2]      |
|---------------------|------------------|

我现在这样做的方式涉及一个嵌套的 for 循环，该循环遍历 df1 中的每一行，df2 中的每一行，并找到两个列表之间的交集，然后将其附加到一个空列表中。

代码如下：

# initialize two empty lists, one for each column
col_a_intersect = []
col_b_intersect = []

# loop over both dataframes to find intersection of lists row by row
for index_1, row_1 in df1.iterrows():
  for index_2, row_2 in df2.iterrows():
    col_a_intersect.append(list(set(row_1['col_a'])& set(row_2['col_a'])))
    col_b_intersect.append(list(set(row_1['col_b'])& set(row_2['col_b'])))

# create a dataframe of the two lists that are now populated
output = pd.DataFrame(
  {col_a: col_a_intersect,
   col_b: col_b_intersect})

代码按预期工作，但是运行需要非常非常长的时间。我正在处理 ~20K x ~100K 的数据集，现在已经运行了大约 18 个小时。 CPU 和 RAM 的使用可以忽略不计，因此任务本身并不是非常占用资源。我怀疑长时间的持续时间是由于使用了 for 循环，这通常是不受欢迎的。有没有更快的方法来完成我所描述的任务？

【问题讨论】：

标签： python pandas dataframe loops optimization

【解决方案1】：

使用这种类型的数据，您不会获得太多矢量化。也就是说，您可以分解两列并进行合并，然后是 groupby：

(df1.reset_index()
    .explode('col_a').explode('col_b')
    .merge(df2.explode('col_a').explode('col_b'), on=['col_a','col_b'])
    .groupby(['index'])
    .agg(lambda x: list(x.unique()))
)

输出：

          col_a col_b
index                
0      [A2, A4]  [B1]
1      [A2, A4]  [B2]
2          [A5]  [B2]

【讨论】：

知道了。我想避免在任何步骤中将其炸毁的原因是因为生成的数据帧是非常大的行（约 40M 行），它通过机器中的 RAM。