Creating conditional columns on Pandas with Numpy select() and where() methods | by B. Chen | Towards Data Science

Pandas is an amazing library that contains extensive built-in functions for manipulating data. Although the built-in functions are capable of performing efficient data analysis, incorporating methods from other library adds value to Pandas.
In this article, we are going to take a look at how to create conditional columns on Pandas with Numpy select() and where() methods
Please check out my Github repo for the source code
Suppose we have a dataset about a fruit store.
df = pd.DataFrame({    'fruit': ['Apple', 'Banana', 'Apple', 'Banana'],    'supplier': ['T & C Bro', 'T & C Bro', 'JM Wholesales', 'JM Wholesales'],    'weight (kg)': [1000,2000,3000,4000],    'customer_rating': [4.8,3.2, 4.2, 4.3]})
We can see that the store has purchased fruits from 2 suppliers “T & C Bro” and “JM Wholesales”. Their price lists are kept separately like below:
# T & C Brotc_price = pd.DataFrame({    'fruit': ['Apple', 'Banana', 'Orange', 'Pineapple'],    'price (kg)': [1.1, 2, 2.9, 3.1]})# JM Wholesalesjm_price = pd.DataFrame({    'fruit': ['Apple', 'Banana', 'Orange', 'Pineapple'],    'price (kg)': [1.2, 1.8, 4, 6]})
We would like to retrieve the price from the supplier’s price list, combine them and add the result to a new column. The expected output is:
The tricky part in this calculation is that we need to retrieve the price (kg) conditionally (based on supplier and fruit) and then combine it back into the fruit store dataset.
For this example, a game-changer solution is to incorporate with the Numpy where() function. A single line of code can solve the retrieve and combine.
df = df.set_index('fruit')tc_price = tc_price.set_index('fruit')jm_price = jm_price.set_index('fruit')
By calling set_index('fruit') to set the fruit column as index across all datasets. This is important so we can use loc[df.index] later to select a column for value mapping.
The Numpy where(condition, x, y) method [1] returns elements chosen from x or y depending on the condition. The most important thing is that this method can take array-like inputs and returns an array-like output.
df['price (kg)'] = np.where(    df['supplier'] == 'T & C Bro',     tc_price.loc[df.index]['price (kg)'],     jm_price.loc[df.index]['price (kg)'])
By executing the above statement, you should get the output like below:
A bit confusion? Here are some details:
df['supplier'] == 'T & C Bro' returns a boolean array
df.index returns ['Apple', 'Banana', 'Apple', 'Banana'] (Index set by Step 1). And both tc_price.loc[df.index] and jm_price.loc[df.index] return a same length DataFrame based on label df.index.
We have learnt how to create a conditional column from 2 datasets. What about more than 2 datasets, for example, 3 different suppliers in the fruit store dataset.
For more than 2 datasets/choices, we can use Numpy select() method instead. Let’s see how it works with the help of an example.
The first step is to combine all price lists into one dataset.
After that, set the fruit column as index.
df_3 = df_3.set_index('fruit')df_price = df_price.set_index('fruit')
Numpy select(condlist, choicelist) method returns an array drawn from elements in choicelist, depending on condition in condlist.
args = df_price.loc[df_3.index]conds = [    df_3['supplier'] == 'T & C Bro',     df_3['supplier'] == 'JM Wholesales',     df_3['supplier'] == 'Star Ltd.',]choices = [    args['T & C Bro'],     args['JM Wholesales'],     args['Star Ltd.'],]df_3['price (kg)'] = np.select(conds, choices)
Basically,
If the condition df_3['supplier'] == 'T & C Bro' is satisfied, it takes the ouput elements from args['T & C Bro'] .
If the condition df_3['supplier'] == 'JM Wholesale' is satisfied, it takes the ouput elements from args['JM Wholesale'] .
If the condition df_3['supplier'] == 'Star Ltd.' is satisfied, it takes the ouput elements from args['Star Ltd.'] .
By executing the above statement, you should get the output like below:
Thanks for reading.
Please checkout the notebook on my Github for the source code.
Stay tuned if you are interested in the practical aspect of machine learning.
When to use Pandas transform() function
A Practical Introduction to Pandas pivot_table() function
Using Pandas method chaining to improve code readability
Working with datetime in Pandas DataFrame
Pandas read_csv() tricks you should know
4 tricks you should know to parse date columns with Pandas read_csv()
More can be found from my Github
[1] Numpy Documentation — where()
[2] Numpy Documentation — select()