如何提高熊猫的迭代操作速度

时间:2018-07-30 08:59:30

标签: python performance pandas

我有一个数据框,其中包含成千上万个条目,其中包含不同变量组合的回归结果。使用单个变量列表和itertools组合函数形成用于回归的组合。

我现在正在寻找一种删除变量组合的方法,这些变量组合代表相似的度量,因此需要删除。我列出了所有不能一起出现的变量。我的代码遍历包含组合的数据框,并使用collections.Counter函数对重复列表中每一行中的元素数进行计数。如果该行中有两个以上的元素,则该行不会复制到新的已清理数据框中。我的代码如下:

import pandas as pd
import itertools
import random
from collections import Counter

def remove_elements(data, dup_col, duplicate_list):
    """
    Removes items from a dataframe that contain multiple items from the duplicate list.

    Arguments:
        data (dataframe): Dataframe containing the data
        dup_col (string): Column name for within data to check for duplicates
        duplicate_list (list): List of duplicate items

    Returns:
        cleaned dataframe
    """
    df_cleaned = pd.DataFrame()
    for idx, row in df.iterrows():
        if any(ele in  duplicate_list for ele in row[dup_col]):
            lenduplicate = sum(Counter(set.intersection(set(duplicate_list), set(row[dup_col]))).values())
            if lenduplicate > 1:
                continue
            else:
                df_cleaned = df_cleaned.append(row)
        else:
            df_cleaned = df_cleaned.append(row)

    return df_cleaned

#%% Create Data

df = pd.DataFrame()
models = ['A', 'B', 'C', 'D', 'E', 'F', 'G']
df['model'] = [list(x) for x in itertools.combinations(models, 4)]
df['result'] = [random.randint(1,101) for i in range(0,len(df['model']))]


# Run Function
same_elements = ['A', 'D', 'G'] 
df = remove_elements(df, 'model', same_elements)

该函数似乎可以正常工作,但是在包含数千个条目的数据帧上,它占用20到30分钟。是否可以使此操作更快?

欢迎任何建议。

BJR

1 个答案:

答案 0 :(得分:1)

您应该能够使用np.intersect1d创建一个掩码,该掩码的交点长度为两个或更多,并对其进行布尔翻转以仅保留原始数据帧中所需的元素。

df[~df.model.apply(lambda L: len(np.intersect1d(L, same_elements)) >= 2)]

将根据您现有的代码为您提供相同的输出:

           model  result
1   [A, B, C, E]      39
2   [A, B, C, F]      62
7   [A, B, E, F]      10
13  [A, C, E, F]      28
20  [B, C, D, E]      38
21  [B, C, D, F]      33
24  [B, C, E, G]       9
25  [B, C, F, G]      11
26  [B, D, E, F]      73
29  [B, E, F, G]       1
30  [C, D, E, F]      96
33  [C, E, F, G]      77