How to apply a function to two columns of Pandas dataframe

Suppose I have a function and a dataframe defined as below:

def get_sublist(sta, end):
    return mylist[sta:end+1]

df = pd.DataFrame({'ID':['1','2','3'], 'col_1': [0,2,3], 'col_2':[1,4,5]})
mylist = ['a','b','c','d','e','f']

Now I want to apply get_sublist to df's two columns 'col_1', 'col_2' to element-wise calculate a new column 'col_3' to get an output that looks like:

  ID  col_1  col_2            col_3
0  1      0      1       ['a', 'b']
1  2      2      4  ['c', 'd', 'e']
2  3      3      5  ['d', 'e', 'f']

I tried:

df['col_3'] = df[['col_1','col_2']].apply(get_sublist, axis=1)

but this results in the following: error:

TypeError: get_sublist() missing 1 required positional argument:

How do I do this?

← Назад к списку