2
$\begingroup$

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'},
                             {'breakfast'}]
                    }).set_index('day')

Does anyone know how to do this? Thanks!

Improve this question
$\endgroup$
4
  • 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
    Commented Aug 15, 2019 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
    Commented Aug 15, 2019 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$
    – Fatemeh Asgarinejad
    Commented Aug 17, 2019 at 20:29
  • $\begingroup$ @FatemehAsgarinejad The goal dataframe was provided and so is definitely what the question wanted. $\endgroup$
    – myseun
    Commented Aug 17, 2019 at 22:08

1 Answer 1

Reset to default
1
$\begingroup$

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 

df.groupby(['day','meal']).count().reset_index()\ 
  .merge(df.groupby('day')['person'].apply(lambda x: len(set(x))).reset_index(), 
         on=['day','person'])\ 
  .groupby(['day'])['meal'].apply(set)

the final df

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

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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