Home > Software design >  Generate pandas DataFrame column from list of lists of indexes using list comprehension
Generate pandas DataFrame column from list of lists of indexes using list comprehension

Time:12-19

Here is the function I wrote to apply to each row of a df

def use_list_of_list_of_indexes(column1, column2):
"""
Column that contains list of lists of indices is used to generate a new column based on indices
_
Args:
column1 - list of elements
column2 - list of lists of indexes
_
Returns:
pandas.core.series.Series
"""
list_of_lists = column1.apply(lambda row: [[row[i] for i in x] for x in column2])
return list_of_lists

Here is df

col1                     col2
[55, 475,243,60]         [[0,1,2], [0,1,3], [0,2,3], [1,2,3]]

When using

df['col3'] = use_list_of_list_of_indexes(df.col1, df.col2)

I want to get

col3
[[55, 475, 243], [55, 475, 60], [55, 243, 60], [475, 243, 60]]

However, I get error

TypeError: list indices must be integers or slices, not list

Adding [0] to df.col2 solves the problem, but how should I rewrite use_list_of_list_of_indexes function to be able to apply it to the whole column?

df['col3'] = use_list_of_list_of_indexes(df.col1, df.col2[0])

CodePudding user response:

If you have to do it this way you are missing one loop:

column1.apply(lambda row: [[[row[z] for z in y] for y in x] for x in column2])

x gives you [[0,1,2], [0,1,3], [0,2,3], [1,2,3]]

y gives you [0,1,2]

so you have to use z to iterate over y

To do it in a slightly better way:

df['col3'] = df[['col1', 'col2']].apply(lambda r: [[r['col1'][y] for y in x] for x in r['col2']], axis = 1)
  • Related