lazy_frog
lazy_frog

Reputation: 165

pandas dataframe select list value from another column

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!

Upvotes: 2

Views: 54

Answers (1)

jezrael
jezrael

Reputation: 862541

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

Upvotes: 3

Related Questions