Extracting element from a list in a pandas column based on the value in another column

Question

I have a following pandas df and I would like to extract the element from the list column based on whatever number that is on the num column:

list             num
[1,2,3,4,5]       3
[7,2,1,3,4]       4

To obtain:

list             num    element
[1,2,3,4,5]       3        4
[7,2,1,3,4]       4        4

I have tried:

df['element'] = df['list'].apply(lambda x: x[df['num'].apply(lambda y: y)])

But I got TypeError: list indices must be integers or slices, not Series .

Is there anyway I can do this? Thank you!

Answer 1

Use DataFrame.apply per rows with axis=1 :

df['element'] = df.apply(lambda x: x['list'][x['num']], axis=1)
print (df)
              list  num  element
0  [1, 2, 3, 4, 5]    3        4
1  [7, 2, 1, 3, 4]    4        4

Or list comprehension with zip :

df['element'] = [x[y]  for x, y in zip(df['list'], df['num'])]

If possible some values not match of list, here is possible use:

def func(a, b):
    try:
        return a[b]
    except Exception:
        return np.nan    

df['element'] = df.apply(lambda x: func(x['list'], x['num']), axis=1)

Answer 2

Using numpy fancy index

list_val = np.array(df.list.values.tolist())
num_val = df.num.values
df['element'] = list_val[np.arange(df.shape[0]), num_val]

Out[295]:
              list  num  element
0  [1, 2, 3, 4, 5]    3        4
1  [7, 2, 1, 3, 4]    4        4