Conditional column without apply

Pandas data engineering interview problem. Difficulty: beginner. Pattern: Filtering. About 10 minutes. Free to practice.

Create a price tier column from rules in a vectorized way, handling missing prices. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

Adding a column that depends on conditions is a very common job. In pandas, the fast way is a vectorized expression that works on the whole column at once, not a function that is called for each row (apply). df has id and price, and price can be missing. Add a tier column: budget when the price is below 10, standard when it is at least 10 and below 50, premium when it is 50 or more, and unknown when the price is missing. Return id, price and tier, keep the row order, and reset the index. The editor builds df (using np.nan for the missing price) and ends with result = df. Change the last line. np.select is a good fit.

Requirements

• Add tier with the values budget, standard, premium and unknown.
• A missing price gives unknown.
• Use a vectorized expression such as np.select, not apply.
• Return id, price, tier, in the original order, with a clean index, in result.

Constraints

• The border values: 10 is standard and 50 is premium.
• Check for a missing price first, because any comparison with NaN is False.
• np.select takes a list of conditions and a list of choices, and checks them in order.

Examples

Input: df id | price 1 | 4.99 2 | 10 3 | 49.99 4 | 50 5 | NULL 6 | 0 Output: id | price | tier 1 | 4.99 | budget 2 | 10 | standard 3 | 49.99 | standard 4 | 50 | premium 5 | NULL | unknown 6 | 0 | budget 10.0 is standard (the lower boundary is inclusive) and 50.0 is premium. The missing price is unknown.

Topics: lakebench, pandas, np.select, vectorized, apply.

More interview problems · All interview problems · Learn data engineering