CODEWITHSUNDEEP
pythonpandasnumpy
Jivan
Jivan

Reputation: 23098

Select a column based on another column's value with Pandas

With a Pandas dataframe structured as follows:

              color    price_blue   price_red   price_green
date
2019-01-01    green            10          20            30
2019-01-02      red            40          50            60
2019-01-03     blue            11          12            13

How can we add a price column based on the value of the color column for each row:

              color    price_blue   price_red   price_green    price
date
2019-01-01    green            10          20            30       30
2019-01-02      red            40          50            60       50
2019-01-03     blue            11          12            13       11

Although an apply-like solution seems easy enough, in this case a vectorised solution is needed as the dataset is relatively large.

In addition, solutions like np.where work fine with only a few color choices, however in the real case there are more than 20 different "colors", so that would end up quite impractical with a lot of nested where.

Upvotes: 1

Views: 486

Answers (1)

jezrael
jezrael

Reputation: 863731

Use DataFrame.lookup with add price_ to column color for match by columns names:

df['price'] = df.lookup(df.index, 'price_' + df['color'])
print (df)
            color  price_blue  price_red  price_green  price
2019-01-01  green          10         20           30     30
2019-01-02    red          40         50           60     50
2019-01-03   blue          11         12           13     11

Upvotes: 5

Related Questions