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


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(
    lambda x: x.nlargest(3).sum()).nlargest(
  • 1
    \$\begingroup\$ You should explain a little bit more about the function you're suggesting. \$\endgroup\$ – Mario Santini Jan 3 '17 at 12:26

SUMLEV is explained here

Definitely want to use nlargest
The advantage of nlargest is that it performs a partial sort therefore in linear time.

However, you don't want to groupby twice. So we'll define a helper function in which we only groupby once.

I'm using a lot of .values to access underlying numpy objects. This will improve some efficiency a bit.

Lastly, I don't like accessing names from the outer scope without a purpose so I'm passing the dataframe as a parameter.

def answer_six(df):
    # subset df to things I care about
    sumlev = df.SUMLEV.values == 50
    data = df[['CENSUS2010POP', 'STNAME', 'CTYNAME']].values[sumlev]

    # build a pandas series with State and County in the index
    # vaues are from CENSUS2010POP
    s = pd.Series(data[:, 0], [data[:, 1], data[:, 2]], dtype=np.int64)

    # define a function that does the nlargest and sum in one
    # otherwise you'd have to do a second groupby
    def sum_largest(x, n=3):
        return x.nlargest(n).sum()

    return s.groupby(level=0).apply(sum_largest).nlargest(3).index.tolist()



['California', 'Texas', 'Illinois']
census_df[census_df['SUMLEV'] == 50].groupby(
    lambda x: x.nlargest(3).sum()).nlargest(

Above: This seems to be the best way to do it. Below: I found another way that is slightly less elegant, but it helped me understand why sgDysregulation's solution works. I hope it can help you also.

census_df[census_df['SUMLEV'] == 50].groupby(
  • \$\begingroup\$ Welcome to Code Review! You have presented an alternative solution, but haven't reviewed the code. Please edit it to explain your reasoning (how your solution works and how it improves upon the original) so that everyone can learn from your thought process. \$\endgroup\$ – Toby Speight Sep 21 '17 at 8:55

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.