如何使数据帧更快?用字典还是用numpy?

Question

我是数据结构的新手，我想让我的代码更快(这只是更大代码的一部分)。在查找变量时使用数据框架会减慢获取实例的速度。我有两个daraframes

df = pd.DataFrame({ 'id1':['A', 'B', 'C','D','J','K','H'], 'id2':[['K','J'], 'NaN',['A'],['B','C'],'NaN','NaN',['A']],'length':[2,'NaN',1,2,'NaN','NaN',1]})

其中length为列表(id2)中元素的个数，

df2 = pd.DataFrame({ 'col1':['B', 'C', 'A','K','J','A'], 'col2':['D', 'D', 'C','A','A','H'],'type':['FF', 'SS', 'FS','SF','FS','FF']}) 

慢的部分是

import time
startTime = time.time()
key='D'
if df.loc[df['id1'] == key, 'id2'].any():
print('continue')
for pred in df[df['id1']==key]['id2'].values[0]:
reltype=df2.loc[(df2['col1'] == pred) & (df2['col2']==key), 'type'].values[0]
print(reltype)
executionTime = (time.time() - startTime)
print('Execution time in seconds: ' + str(executionTime))

我所做的是将df转换为字典并重写我的代码如下，但我不知道如何重写'reltype'部分并使其更快。

df_dic=df.groupby('id1')[['id2', 'length']].apply(lambda g: g.values.tolist()).to_dict()
key='D'
if ~np.isnan(df_dic[key][0][1]):
print('continue')
for pred in df_dic[key][0][0]:
reltype=df2.loc[(df2['col1'] == pred) & (df2['col2']==key), 'type'].values[0]
print(reltype)

我有兴趣使上面的代码更快，特别是在寻找'reltype'的搜索。我将感谢所有的建议。由于

Answer 1

我认为你的第一个实现应该很快。如果不重复切片，可能会有一个小小的改进:

for key in ['D']:
s = df.loc[df['id1'] == key, 'id2']
if s.any():
print('continue')
df_temp = df2.loc[df2['col2'] == key, :]
for pred in s.values[0]:
if pred != NaN:
reltype = df_temp.loc[df_temp['col1'] == pred, 'type'].values[0]
print(reltype)

根据你的数据使用df_temp可能比使用布尔掩码进行切片

慢。