我怎樣才能改善這段代碼的運行時間？

Question

我想在兩個數據框之間分享一些信息。 我的代碼有效，但需要很長時間。 你知道我可以如何改善我的運行時間嗎？ 我正在嘗試執行以下操作：

我有一個 dataframe df1 （它有 160 列，但這里只有顯示的很重要）：

         a_idx  b_idx  c_idx  d_idx  e_idx  f_idx Evt_ID
    0    0      1      3      4      2      6     346642
    1    1      2      3      4      5      5     917426
    2    0      1      3      4      2      2     123543
                        ...
                        

還有一個 dataframe df2 （ist 有 10 列，但這里只有這些很重要）：

    Name    Evt_ID
0   Jet1    346642
1   Jet2    346642
2   Jet3    346642
3   Jet4    346642
4   Jet5    346642
5   Jet6    346642
6   Jet7    346642
7   Lepton  346642
8   Jet1    917426
9   Jet2    917426
      ...

現在我想在df2中創建一個名為“y”的新列，其中包含每行的類別。 類別可以在df1的幫助下找到，類別是： category_list = ["a", "b", "c", "d", "e", "f"]也可以是"unknown" 。 例如，df1 中的第一行的值category = [0,1,3,4,2,6]這意味着df2應該如下所示：

（解釋： category中的第五個數字是2 --> Jet( 2 +1) = Jet3 在 category_list 中有第五個category_list ： “e” ）

    Name    Evt_ID    y
0   Jet1    346642    a
1   Jet2    346642    b
2   Jet3    346642    e
3   Jet4    346642    c
4   Jet5    346642    d
5   Jet6    346642    unknown
6   Jet7    346642    f
7   Lepton  346642    unknown
     ...

我實現這一目標的方法如下：

df["y"] = "unknown"
category_list = ["a", "b", "c", "d", "e", "f"]

for event_id in tqdm(df1.Evt_ID):
    category = df1.loc[df1.Evt_ID == event_id, ["a_idx","b_idx",
                                               "c_idx", "d_idx", 
                                               "e_idx", "f_idx"]].values.squeeze()
    
    i = 0
    for jet_index in category:
        df2.loc[(dfo.Evt_ID == event_id) & (dfo.Name == "Jet".join(str(jet_index+1))), "y"] = category_list[i] 
        i += 1

此代碼需要 30 或 60 分鍾才能運行，具體取決於運行它的 jupyter notebook。為什么 notebook 本身會影響運行時？ 但更重要的是：如何改進運行時間？

Answer 1

由於其矢量化結構，以下代碼段應該運行得更快。

這里有兩個技巧。 第一個是使用df.melt有效地將列a 、 b 、 ...、 e轉換為行。 第二個是join生成的 DataFrame 與df2連接起來。 這樣，所有缺失值都變為NaN並且可以用df.fillna替換為unknown 。

cols = ["a_idx", "b_idx", "c_idx", "d_idx", "e_idx", "f_idx"]
df = df1[cols + ["Evt_ID"]].rename(columns={c: c[0] for c in cols})

df = df.melt(id_vars="Evt_ID", var_name="y")
df["value"] = "Jet" + (df["value"] + 1).astype(str)

df = df2.join(df.set_index(["Evt_ID", "value"]), on=["Evt_ID", "Name"])
df = df.fillna("unknown")

最后， df看起來像：

     Name  Evt_ID        y
0    Jet1  346642        a
1    Jet2  346642        b
2    Jet3  346642        e
3    Jet4  346642        c
4    Jet5  346642        d
5    Jet6  346642  unknown
6    Jet7  346642        f
7  Lepton  346642  unknown
8    Jet1  917426  unknown
9    Jet2  917426        a

該結果是通過以下樣本數據獲得的：

import pandas as pd


df1 = pd.DataFrame(
    [
        [0, 1, 3, 4, 2, 6, 346642],
        [1, 2, 3, 4, 5, 5, 917426],
        [0, 1, 3, 4, 2, 2, 123543],
    ],
    columns=["a_idx", "b_idx", "c_idx", "d_idx", "e_idx", "f_idx", "Evt_ID"],
)

df2 = pd.DataFrame(
    [
        ["Jet1", 346642],
        ["Jet2", 346642],
        ["Jet3", 346642],
        ["Jet4", 346642],
        ["Jet5", 346642],
        ["Jet6", 346642],
        ["Jet7", 346642],
        ["Lepton", 346642],
        ["Jet1", 917426],
        ["Jet2", 917426],
    ],
    columns=["Name", "Evt_ID"],
)