python - 优化大熊猫数据框的重复查询（掩码）

Question

我有一个熊猫数据框，最初有大约 350 列和 500000 行：

import string
import numpy as np
import pandas as pd
import itertools

cols = list(string.ascii_lowercase) + [i+j for i,j in [*itertools.combinations(list(string.ascii_lowercase), 2)]]
df = pd.DataFrame({col: np.repeat([np.random.randint(2)], [500000]) for col in cols})

我需要向我的数据框添加 3000 个新列（初始化为 0），其中每行的值取决于现有行中的值（我在测试时使用掩码）：

for i, j, k in itertools.combinations(list(string.ascii_lowercase), 3):
    df[i+j+k] = 0
    df.loc[(df[i] > 0) & (df[j] > 0) & (df[k] > 0) & (df[i + j] + df[i + k] + df[j + k] >= 2), i+j+k] = 1

但是，问题是上述循环非常慢！有没有办法优化上述程序？也许使用更快的熊猫查找功能？

标签： pythonpandasoptimization

Answer 1

在这里，您有一个更快的解决方案。当您遇到此类性能问题时，请尝试切换到 Numpy。它的速度更快！

在运行之前确保它适合您的记忆：D

import string
import numpy as np
import pandas as pd
import itertools

cols = list(string.ascii_lowercase) + [i+j for i,j in [*itertools.combinations(list(string.ascii_lowercase), 2)]]
df = pd.DataFrame({col: np.repeat([np.random.randint(2)], [500000]) for col in cols})
mat = df.values  # Convert to numpy, much faster


additional_cols = []
for i, j, k in itertools.combinations(list(string.ascii_lowercase), 3):
    cond_1 = (mat[:,cols.index(i)]*mat[:,cols.index(j)]*mat[:,cols.index(k)])>0  # Singles
    cond_2 = (mat[:,cols.index(i+j)]+mat[:,cols.index(i+k)] + mat[:,cols.index(j+k)])>2  # Multiples
    col = (cond_1 & cond_2) + 0  # +0 transforms to int
    additional_cols.append((i+j+k, col))

df_additional = pd.DataFrame(dict(additional_cols)) # Assure it fits in memory
df = pd.concat([df, df_additional], axis=1)  # Assure it fits in memory