Select column dynamically in Pandas dataframe based on values in a list or another column

Question

I have a pandas data frame that looks like this -

For each row I want to select value from either Col1 or Col2 based on value in INDX. So the output in above case should be- [10,40,60,70]

I did this by looping through each row of dataframe, but it's quite slow. Is there is a faster way to accomplish this?

Dummy test code -

for i in np.arange(0, df.shape[0]):
    print(df.iloc[i, df['INDX'][i]])

Quang Hoang · Accepted Answer · 2021-01-25 20:58:44Z

4

Try lookup:

cols = df.columns[:2]

df.lookup(df.index, cols[df.INDX])

Output:

array([10, 40, 60, 70])

Update As commented by Scott, lookup is deprecated. We can resolve to numpy indexing:

df[cols].to_numpy()[np.arange(len(df)), df['INDX']]

answered Jan 25, 2021 at 20:54

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Read through these issues on pandasdev github. github.com/pandas-dev/pandas/issues/39171

Can't help putting a comment in the issue on Github.