Pandas 数据框按 3 列和条件搜索

Question

我有一个来自美国的人口普查数据的 pandas dataframe。列包括STNAME（州名）CTYNAME（县名）以及按年份在人口普查报告中收集的不同数据的许多列。

我正在寻找人口最多的 3 个州，仅包括每个州人口最多的 3 个县。

这是我当前的代码，它稍微整理了原始数据帧，只包含与问题相关的内容：

def answer():
    census_df50 = census_df[census_df['SUMLEV'] == 50]
    columns = ['STNAME', 'CTYNAME', 'CENSUS2010POP']
    c = census_df50[columns]
    return c
print(answer())

下面是终端打印的数据框示例：

     STNAME             CTYNAME  CENSUS2010POP
1       Alabama      Autauga County          54571
2       Alabama      Baldwin County         182265
3       Alabama      Barbour County          27457
4       Alabama         Bibb County          22915
5       Alabama       Blount County          57322
6       Alabama      Bullock County          10914
7       Alabama       Butler County          20947
8       Alabama      Calhoun County         118572

它是按州和县的字母顺序列出的，所以这只按县名显示阿拉巴马州的人口普查数据，但数据框中有 3000 多行（每个县一个，州有多个条目）

我的方法是编写一个函数，根据 CENSUS2016POP 的值（代表 2010 年人口的整数）在 CTYNAME 中查找前 3 个县，为 STNAME 列中的每个州。然后让函数按照此标准将名称作为前 3 个状态的字符串返回。但是，对于如何实现这一点非常迷茫。

我很确定我应该尝试使用.groupby()、.set_index() 和 .nlargest() 这些函数的组合。

任何帮助将不胜感激！

Answer 1

由于您没有提供任何示例数据，以下是一些：

STNAME,CTYNAME,POPULATION
A,A1,100
A,A2,20
A,A3,30
A,A4,40
B,B1,10
B,B2,2
B,B3,30
B,B4,40
C,C1,100
C,C2,20
C,C3,300
C,C4,40
D,D1,10
D,D2,20
D,D3,30
D,D4,40

In [1]: df = pd.read_clipboard(sep=',')

这是达到预期结果的一种方法，也许有更简单但我无法进一步减少它：

In [2]: df.ix[df['STNAME'].isin(df.groupby('STNAME')['POPULATION'].sum().nlargest(3).index)].groupby(['STNAME','CTYNAME']).sum()['POPULATION'].groupby(level=0, group_keys=False).nlargest(3)

---

分解：

# Find the first STNAME Groups
In [3]: largest_states = df.groupby('STNAME')['POPULATION'].sum().nlargest(3).index
        largest_states

Out[3]: Index(['C', 'A', 'D'], dtype='object', name='STNAME')

过滤这些上的df：

In [4]: df2 = df.ix[df['STNAME'].isin(df.groupby('STNAME')['POPULATION'].sum().nlargest(3).index)]

然后在里面找到前三个：

In [5]: df2.groupby(['STNAME','CTYNAME']).sum()['POPULATION'].groupby(level=0, group_keys=False).nlargest(3)

Out[5]:
STNAME  CTYNAME
A       A1         100
        A4          40
        A3          30
C       C3         300
        C1         100
        C4          40
D       D4          40
        D3          30
        D2          20
Name: POPULATION, dtype: int64

---

更新：

如果要按州人口排序，请按以下顺序进行：

In[6]: df2 = df.groupby(['STNAME','CTYNAME']).sum()['POPULATION'].groupby(level=0, group_keys=False).nlargest(3)

       df2.ix[df.groupby('STNAME')['POPULATION'].sum().nlargest(3).index]

Out[6]:
STNAME  CTYNAME
C       C3         300
        C1         100
        C4          40
A       A1         100
        A4          40
        A3          30
D       D4          40
        D3          30
        D2          20
Name: POPULATION, dtype: int64