How to extract element of pandas list column specified in other column?

Question

I have a Pandas dataframe with two columns:

• col1: a list column
• col2: an integer that specifies the index of the list element that I would like to extract and store in col3. It can take NaN value, in which case the outcome should be NaN as well.

Sample input:

df = pd.DataFrame({
    'col1' : [['A', 'B'], ['C', 'D', 'E'], ['F', 'G']], 
    'col2' : [0, 2, np.nan]})

Expected output:

df_out = pd.DataFrame({
    'col1' : [['A', 'B'], ['C', 'D', 'E'], ['F', 'G']], 
    'col2' : [0, 2, np.nan],
    'col3' : ['A', 'E', np.nan]})

Answer 1

You can do a list comprehension to compare the two columns too.

df['col3'] = [i[int(j)] if not np.isnan(j) else j for i,j in zip(df.col1,df.col2)]

        col1  col2 col3
0     [A, B]   0.0    A
1  [C, D, E]   2.0    E
2     [F, G]   NaN  NaN

Answer 2

you can use a basic apply:

def func(row):
    if np.isnan(row.col2):
        return np.nan
    else:
        return row.col1[int(row.col2)]


df['col3'] = df.apply(func, axis=1)

output:

        col1  col2 col3
0     [A, B]   0.0    A
1  [C, D, E]   2.0    E
2     [F, G]   NaN  NaN