Python - pick a value from a list basing on another list

Question

I've got a dataframe. In column A there is a list of integers, in column B - an integer. I want to pick n-th value of the column A list, where n is a number from column B . So if in columns A there is [1,5,6,3,4] and in column B : 2, I want to get '6'.

I tried this:

result = [y[x] for y in df['A'] for x in df['B']

But it doesn't work. Please help.

Answer 1

You can go for apply ie

df = pd.DataFrame({'A':[[1,2,3,4,5],[1,2,3,4]],'B':[1,2]})


                A  B
0  [1, 2, 3, 4, 5]  1
1     [1, 2, 3, 4]  2

# df.apply(lambda x : np.array(x['A'])[x['B']],1)
# You dont need np.array here, use it when the column B is also a list.
df.apply(lambda x : x['A'][x['B']],1) # Thanks @Zero 
0    2
1    3
dtype: int64

Answer 2

Use zip with list comprehension :

df['new'] = [y[x] for x, y in zip(df['B'], df['A'])]
print (df)
                 A  B  new
0  [1, 2, 3, 4, 5]  1    2
1     [1, 2, 3, 4]  2    3