遍歷 pandas dataframe 並使用自定義 function 創建新列

Question

我有一個 pandas dataframe （顯然）包含一些數據。 我創建了一個輸出一些新列的 function。 如何迭代或應用 function？

我在下面創建了一個最小示例（不是實際問題），帶有 dataframe 和 function。

編輯：將 function 視為“黑匣子”。 我們現在不知道里面有什么，但根據輸入它返回一個 dataframe，它應該添加到現有的 dataframe 中。

import pandas as pd
a=pd.DataFrame({"input1": ["a","b"], "input2":[3,2]})

  input1  input2
0      a       3
1      b       2

def f(i1,i2):
    return(pd.DataFrame([{"repeat" : [i1]*i2, "square":i2**2 }]))

所以在這種情況下，function 返回兩個新列“repeat”和“square”

f(a.iloc[0,0],a.iloc[0,1])

      repeat  square
0  [a, a, a]       9

f(a.iloc[1,0],a.iloc[1,1])
   repeat  square
0  [b, b]       4

我想以這樣的數據框結束

  input1  input2     repeat  square
0      a       3  [a, a, a]       9
1      b       2     [b, b]       4

有沒有人對此有一個優雅的解決方案？

Answer 1

我會使用assign來做到這一點：

a = a.assign(
    repeat = a['input1'].repeat(a['input2']).groupby(level=0).agg(list),
    square = np.square(a['input2']),
)

Output：

>>> a
  input1  input2     repeat  square
0      a       3  [a, a, a]       9
1      b       2     [b, b]       4

Answer 2

您可以嘗試對f function 進行此修改：

import pandas as pd


def f(df, col1, col2):
    df_ = pd.DataFrame([{"repeat": list(df[col1] * df[col2])}]).explode("repeat", ignore_index=True)
    df_["square"] = list(df[col2] ** 2)
    return pd.concat([df, df_], axis=1)


a = pd.DataFrame({"input1": ["a", "b"], "input2": [3, 2]})
f(a, "input1", "input2")

Answer 3

使用pd.concat怎么樣？

generated_df = pd.concat([f(*args) for args in a.to_numpy()], ignore_index=True)
out = pd.concat([a, generated_df], axis=1)

Output：

>>> out
  input1  input2     repeat  square
0      a       3  [a, a, a]       9
1      b       2     [b, b]       4