Finding the states with the three most populous counties I just started to use Python and Pandas. My current solution to a problem looks ugly and inefficient. I would like to know how to improve it. Data file is Census 2010 can be viewed here

Question:


  Only looking at the three most populous counties for each state, what are the three most populous states (in order of highest population to lowest population)? Use CENSUS2010POP.
  This function should return a list of string values.


Current code

def answer_six():
    cdf = census_df[census_df['SUMLEV'] == 50]
    cdf = cdf.groupby('STNAME')
    cdf = cdf.apply(lambda x:x.sort_values('CENSUS2010POP', ascending=False)).reset_index(drop=True)
    cdf = cdf.groupby('STNAME').head(3)
    cdf = cdf.groupby('STNAME').sum()
    cdf = cdf.sort_values('CENSUS2010POP', axis=0, ascending=False).head(3)
    return list(cdf.index)
nlargest could help, it finds the maximum n values in pandas series. 

In this line of code, groupby groups the frame according to state name, then apply finds the 3 largest values in column CENSUS2010POP and sums them up. The resulting series has the unique state names as index and the corresponding top 3 counties sum, applying nlargest gets a series with the top 3 states as required 

return census_df[census_df['SUMLEV'] == 50].groupby(
    'STNAME')['CENSUS2010POP'].apply(
    lambda x: x.nlargest(3).sum()).nlargest(
    3).index.values.tolist()