Finding the states with the three most populous counties

Question

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)

Answer 1

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()

Answer 2

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()

Demonstration

answer_six(census_df)

['California', 'Texas', 'Illinois']

Answer 3

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

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(
    'STNAME')['CENSUS2010POP'].nlargest(3).groupby(
    'STNAME').sum().nlargest(3).index.values.tolist()