Pandas GroupBy分组返回元组而非值的问题排查与解决

问题描述 #

我有如下Pandas DataFrame：

id_country  txt_template_1  txt_template_2  id_set  id_question     txt_question
0   NEUTRAL     template neutral 1  template neutral 2  1   1   1_1
1   NEUTRAL     template neutral 1  template neutral 2  1   2   1_2
2   NEUTRAL     template neutral 1  template neutral 2  1   3   1_3
3   NEUTRAL     template neutral 1  template neutral 2  1   4   1_4
4   NEUTRAL     template neutral 1  template neutral 2  2   1   2_1
5   NEUTRAL     template neutral 1  template neutral 2  2   2   2_2
6   NEUTRAL     template neutral 1  template neutral 2  2   3   2_3
7   NEUTRAL     template neutral 1  template neutral 2  2   4   2_4
8   FRA     template FRA 1  template FRA 2  1   1   1_1
9   FRA     template FRA 1  template FRA 2  1   2   1_2
10  FRA     template FRA 1  template FRA 2  1   3   1_3
11  FRA     template FRA 1  template FRA 2  1   4   1_4
12  FRA     template FRA 1  template FRA 2  2   1   2_1
13  FRA     template FRA 1  template FRA 2  2   2   2_2
14  FRA     template FRA 1  template FRA 2  2   3   2_3
15  FRA     template FRA 1  template FRA 2  2   4   2_4

我写了嵌套分组的函数，先按id_country分组，再在每个分组内按id_set分组。但循环打印set_id时，输出的是(1,)、(2,)这类元组，而非id_set的单个数值：

def ask_question(df):
  grouped_country = df.groupby(['id_country'])

  # 遍历每个国家分组
  for country_id, group_country_df in grouped_country:
    grouped_id_set = group_country_df.groupby(['id_set'])

    # 遍历每个id_set分组
    for set_id, group_set_df in grouped_id_set:
      print(set_id)

输出结果：

(1,)
(2,)
(1,)
(2,)
(1,)
(2,)

[]

需要知道原因，以及如何让set_id输出具体数值而非元组。

原因分析 #

当你用groupby(['id_set'])（参数是列表形式）时，Pandas为了兼容多列分组的场景，会统一将分组键包装成元组返回——哪怕只有一个分组列，也会返回长度为1的元组，这就是你看到(1,)、(2,)的原因。

解决方法 #

方法1：分组时直接传列名字符串 #

把分组参数从列表改成单个字符串，这样Pandas会直接返回单个数值作为分组键：

def ask_question(df):
  grouped_country = df.groupby(['id_country'])

  for country_id, group_country_df in grouped_country:
    # 直接传列名字符串，而非列表
    grouped_id_set = group_country_df.groupby('id_set')

    for set_id, group_set_df in grouped_id_set:
      print(set_id)

方法2：从元组中提取第一个元素 #

如果需要保留列表形式的分组参数，也可以通过索引[0]取出元组里的具体数值：

def ask_question(df):
  grouped_country = df.groupby(['id_country'])

  for country_id, group_country_df in grouped_country:
    grouped_id_set = group_country_df.groupby(['id_set'])

    for set_id, group_set_df in grouped_id_set:
      # 提取元组的第一个元素
      print(set_id[0])

两种方法的输出都会变成：

1
2
1
2
1
2

内容的提问来源于stack exchange，提问作者Louis