Top N per group

Pandas data engineering interview problem. Difficulty: intermediate. Pattern: Window Functions. About 14 minutes. Part of the Pro drill bank.

Top 2 products by revenue in each category, keeping products that tie for second place. Treat this as a production helper: match the contracted return shape, including empty and duplicate inputs.

You often want the top few items of each group, such as the 2 best products of each category. When two products tie for second place, it is fairer to keep both than to cut one at random. df has category, product and revenue. For each category, return the products with the 2 highest revenues. Products that tie with the second highest revenue are all kept, so a category can return more than 2 rows. A category with fewer than 2 products returns what it has. Return category, product and revenue, sorted by category, then revenue from high to low, then product, with a clean index. The editor builds df and ends with result = df. Change the last line to build your answer.

Requirements

• Rank the revenues inside each category, with ties sharing the same rank.
• Keep the rows with rank 1 or 2.
• Sort by category, revenue descending, then product.
• Reset the index and store the DataFrame in result.

Constraints

• Use groupby("category")["revenue"].rank(method="min", ascending=False) so tied rows get the same rank.
• Do not drop the rank column from the result before you filter on it.
• A category can have more than 2 rows in the result when there are ties.

Examples

Input: df category | product | revenue A | p1 | 100 A | p2 | 90 A | p3 | 90 A | p4 | 50 B | p5 | 70 B | p6 | 30 C | p7 | 10 Output: category | product | revenue A | p1 | 100 A | p2 | 90 A | p3 | 90 B | p5 | 70 B | p6 | 30 C | p7 | 10 Category A returns 3 rows because p2 and p3 tie for second place. B returns both its products and C its only one.

Topics: lakebench, pandas, groupby, rank, ties.

More interview problems · All interview problems · Learn data engineering