在 pandas DataFrame 中有效地搜索列表值的组合

Posted 2023-03-31

【中文标题】在 pandas DataFrame 中有效地搜索列表值的组合

【英文标题】：Efficiently search for combinations of list values in a pandas DataFrame

【发布时间】：2019-05-23 17:22:07

【问题描述】：

我正在迭代多个 pandas 数据帧，在每次迭代中我提取 2 个列表。 然后，对于每 2 个列表组合，我需要在另一个数据框中找到这 2 个列表中所有值的每个组合出现的行。

简化示例：

df 看起来像这样（任何 D_val P_val 可能只出现一次）：

D_val | P_val | D_PCode  
A     | 1     | 25  
D     | 1     | 21  
C     | 2     | 32  
B     | 4     | 35  
D     | 2     | 126  
B     | 1     | 3  
C     | 1     | 312  

对于具有以下列表的单次迭代 -

list1 = [1,2]  
list2 = [A,B,C]  

我希望得到以下列表： [25、3、312、32]

解释：

for [A,1] - 25  
for [A,2] - nothing  
For [B,1] - 3  

等等。

我使用嵌套循环来实现它，但是由于多次迭代、大型数据框和长列表，它需要的时间太长了。

任何建议如何有效地实施它？如果需要，我很乐意使用其他数据结构。

【问题讨论】：

由于您的 D_P 组合是唯一的，您可以为该行（A1、D1、C2 等）创建一个 D 和 P 连接的索引，使用 itertools.product 获得笛卡尔积列表，然后将其用作 df 的掩码。如果这对您有意义，我可以尝试将其扩展为答案。

【参考方案1】：

我们只使用了两次isin

df1=df.loc[df.D_val.isin(list2)&df.P_val.isin(list1)]
df1
Out[211]: 
  D_val  P_val  D_PCode
0     A      1       25
2     C      2       32
5     B      1        3
6     C      1      312

【参考方案2】：

您可以使用itertools.product 生成所有可能的值，然后使用isin：

from itertools import product

import pandas as pd

data = [['A', 1, 25],
        ['D', 1, 21],
        ['C', 2, 32],
        ['B', 4, 35],
        ['D', 2, 126],
        ['B', 1, 3],
        ['C', 1, 312]]

df = pd.DataFrame(data=data, columns=['D_val', 'P_val', 'D_PCode'])

list1 = [1, 2]
list2 = ['A', 'B', 'C']
lookup = set(product(list2, list1))

mask = df[['D_val', 'P_val']].apply(tuple, axis=1).isin(lookup)
result = df[mask].D_PCode

print(result)

输出

0     25
2     32
5      3
6    312
Name: D_PCode, dtype: int64

或者直接在两列上使用 isin，例如：

list1 = [1, 2]
list2 = ['A', 'B', 'C']

result = df[df.D_val.isin(list2) & df.P_val.isin(list1)].D_PCode
print(result)

输出

0     25
2     32
5      3
6    312
Name: D_PCode, dtype: int64

【参考方案3】：

这是一个MultiIndex 问题，最好用reindex 解决。

df = df.set_index(['D_val', 'P_val'])
idx = pd.MultiIndex.from_product([list2, list1])

df.reindex(idx)

     D_PCode
A 1     25.0
  2      NaN
B 1      3.0
  2      NaN
C 1    312.0
  2     32.0

如果您想要匹配的字典，请使用reindex + Series.to_dict：

df.reindex(idx).D_PCode.to_dict()

('A', 1): 25.0,
 ('A', 2): nan,
 ('B', 1): 3.0,
 ('B', 2): nan,
 ('C', 1): 312.0,
 ('C', 2): 32.0

---

如果您只想要存在的行，请使用intersection 来高效地计算交集的索引：

df.loc[df.index.intersection(idx)]

     D_PCode
A 1       25
B 1        3
C 1      312
  2       32