时间戳

首页 > 解决方案 > 在 pandas DataFrame 中有效地搜索列表值的组合

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

问题描述

我正在迭代多个 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

等等

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

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

标签： pythonpandaslistdataframecombinations

解决方案

您可以使用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

推荐阅读