如何使用 apply（基於用戶定義的函數）向數據框添加新行和多行？

Question

對於我的用例，原始框架看起來像 -

指數	col1	col2	col3
0	0	第零個例如	拒絕
1	1	首先例如	接受
2	2	第二個例如	接受
3	3	第三個例如	拒絕

我有一個函數定義為 -

def foo(row):
  if row['col1']==0:
    answers = ['zero']
  elif row['col1']==1:
    answers = ['one', 'i']
  elif row['col1']==2:
    answers = ['two', 'ii']
  else:
    answers = ['three', 'iii']

基於這個函數，我想在我的數據框中添加一個名為 col4 的新列。 本質上，需要添加與answers列表中的值一樣多的新行，其中每行中 col4 的值應該是列表的后續值（而所有其他列的值保持不變）

所以我希望得到的框架像 -

指數	col1	col2	col3	col4
0	0	第零個例如	拒絕	零
1	1	首先例如	接受	一
2	1	首先例如	接受	一世
3	2	第二個例如	接受	二
4	2	第二個例如	接受	ii
5	3	第三個例如	拒絕	三
6	3	第三個例如	拒絕	三

我無法理解我們如何使用 apply 來返回行以及多行。 下面的代碼只會在包含列表的原始框架中添加一個新列 col4 （如果我在foo中返回答案）

input_df['col4'] = input_df.apply(foo, axis=1)

如何修改foo以返回多行？ 任何幫助表示贊賞。

Answer 1

您可以嘗試返回列表然后explode

def foo(row):
  if row['col1']==0:
    answers = ['zero']
  elif row['col1']==1:
    answers = ['one', 'i']
  elif row['col1']==2:
    answers = ['two', 'ii']
  else:
    answers = ['three', 'iii']

  return answers

input_df['col4'] = input_df.apply(foo, axis=1)
input_df = input_df.explode('col4', ignore_index=True)

print(input_df)

   index  col1       col2    col3   col4
0      0     0  zeroth eg  reject   zero
1      1     1   first eg  accept    one
2      1     1   first eg  accept      i
3      2     2  second eg  accept    two
4      2     2  second eg  accept     ii
5      3     3   third eg  reject  three
6      3     3   third eg  reject    iii