I had the following data-cleaning question in an interview test that I struggled on (I've changed the details to anonymise it and protect the company's interview process)

Given the following dataframe df, return a new series with day as the index, and a single column with the set of meals consumed by everyone who ate that day (i.e, both Alice and Bob on days 1 and 3, but only Alice on day 2). Do not use for loops or list comprehensions, only method chaining and a single lambda function that accepts only a single argument.

df = pd.DataFrame({'day':[1, 2, 3, 1, 3]*3,
                   'person':['Alice', 'Alice', 'Alice', 'Bob', 'Bob']*3,
                   'meal':['breakfast', 'breakfast', 'breakfast', 'breakfast', 'breakfast']+
                          ['lunch', 'brunch', 'brunch', 'lunch', 'lunch']+
                          ['dessert', 'dinner', 'snack', 'beer', 'dessert']

In other words, the goal is to obtain the following dataframe:

goal = pd.DataFrame({'day':[1, 2, 3], 
                     'meal':[{'breakfast', 'lunch'}, 
                             {'breakfast', 'brunch', 'dinner'},

Does anyone know how to do this? Thanks!

  • 2
    $\begingroup$ This is kind of a stack overflow question really. Yet, I believe they wanted df.groupby('day').apply(lambda x: set(x['meal'])) $\endgroup$
    – grochmal
    Aug 15 '19 at 18:23
  • $\begingroup$ @grochmal Won’t this give the set of meals eaten by anyone on that day, rather than only those eaten by both? So day 1 = {breakfast, lunch, dessert, beer} instead of {breakfast, lunch}? $\endgroup$
    – myseun
    Aug 15 '19 at 19:40
  • $\begingroup$ @myseun I think the answer by grochmal is what the question wanted, the question wants the whole meals consumed by all the people in each day, whereas in the goal dataset this is not achieved $\endgroup$ Aug 17 '19 at 20:29
  • $\begingroup$ @FatemehAsgarinejad The goal dataframe was provided and so is definitely what the question wanted. $\endgroup$
    – myseun
    Aug 17 '19 at 22:08

The first groupby - counts the number of persons per day per meal

The second groupby - counts the number of unique persons in each day

The inner merge between the 2 - matches the number of persons per day+meal to be equal to the number of persons on that day

We've left with a df that contains day-meal couples in each row, with the number of persons that had a meal which matches the total number of persons on that day:

the new df

Now we only need to groupby day and create the set of the meals

  .merge(df.groupby('day')['person'].apply(lambda x: len(set(x))).reset_index(), 

the final df

  • $\begingroup$ Brilliant, thanks for the explanation with the answer, I've accepted this! $\endgroup$
    – myseun
    Aug 22 '19 at 10:43

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.