pandas dataframe select list value from another column

Question

Everyone! I have a pandas dataframe like this:

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

as we can see, the A column is a list and the B column is an index value. I want to get a C column which is index by B from A:

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

Is there any elegant method to solve this? Thank you!

Answer 1

Use list comprehension with indexing:

df['C'] = [x[y] for x, y in df[['A','B']].to_numpy()]

Or DataFrame.apply , but it should be slowier if large DataFrame:

df['C'] = df.apply(lambda x: x.A[x.B], axis=1)

---

print (df)
           A  B  C
0  [1, 2, 3]  0  1
1  [2, 3, 4]  1  3