迭代数据帧时的嵌套 for 循环优化

时间:2021-05-29 07:04:13

标签: python pandas dataframe function nested-loops

我对 Python 和编码相当陌生。我正在寻找一种优化嵌套 for 循环的方法。 我编写的嵌套 for 循环工作得很好,但运行需要很多时间。 我已经解释了我的原始代码背后的基本思想以及我尝试做的事情,如下:

data = [['a', '35-44', 'male', ['b', 'z', 'x']], ['b', '15-24', 'female', ['a', 'z', 'q']], \
        ['r', '35-44', 'male', ['z', 'a', 'd']], ['q', '15-24', 'female', ['u', 'k', 'b']]]
df = pd.DataFrame(data, columns= ['ID', 'age_group', 'gender', 'matching_ids']) 

df 是我正在处理的 Dataframe。 我想要做的是将 df 中的每个“ID”与同一 df 中的每个其他“ID”进行比较,并检查它是否符合某些条件。

  1. 如果 age_group 相等。
  2. 如果性别相同。
  3. 如果“ID”在“matched_ids”中。

如果满足这些条件,我需要将该行附加到单独的数据帧 (sample_df) 这是带有嵌套 for 循环的代码,可以正常工作:

df_copy = df.copy()
sample_df = pd.DataFrame()
for i in range(len(df)):
    for j in range(len(df)):
        if (i!=j) and (df.iloc[i]['ID'] in df_copy.iloc[j]['matching_ids']) and \
        (df.iloc[i]['gender'] == df_copy.iloc[j]['gender']) and\
        (df.iloc[i]['age_group'] == df_copy.iloc[j]['age_group']):
            sample_df = sample_df.append(df_copy.iloc[[j]])

我尝试通过编写一个函数并使用 df.apply(func) 来简化它,但它仍然需要几乎相同的时间。 下面是使用函数编写的代码:

sample_df_func = pd.DataFrame()
def func_extract(x):
     for k in range(len(df)):
        if (x['ID'] != df_copy.iloc[k]['ID']) and (x['ID'] in df_copy.iloc[k]['matching_ids']) and \
        (x['gender'] == df_copy.iloc[k]['gender']) and\
        (x['age_group'] == df_copy.iloc[k]['age_group']):
            global sample_df_func
            sample_df_func = sample_df_func.append(df_copy.iloc[[k]])
df.apply(func_extract, axis = 1)
sample_df_func

我正在寻找方法来简化并进一步优化它。 请原谅我,如果这个问题的解决方案很简单而我无法弄清楚。

谢谢

PS:两个月前我刚刚开始编码。

1 个答案:

答案 0 :(得分:3)

我们可以在 age_groupgender 上形成组以获得前两个条件自动成立的子集。对于第三个条件,我们可以 explode matching_ids 然后检查 ID 的 any isin ID 并将这些行保留在组中仅使用布尔值索引:

out = (df.groupby(["age_group", "gender"])
         .apply(lambda s: s[s.matching_ids.explode().isin(s.ID).groupby(level=0).any()])
         .reset_index(drop=True))

最后我们重置索引以摆脱分组变量作为索引,

得到

>>> out

  ID age_group  gender matching_ids
0  b     15-24  female    [a, z, q]
1  q     15-24  female    [u, k, b]
2  r     35-44    male    [z, a, d]