Pandas Dataframe，Grouping or Looping？

Question

我的数据框如下：

          Col_A  | Col_B  
        +--------+--------+
Index   
--------+--------+--------+ 
Index_1 | XXXXX  | 0      |
--------+--------+--------+
Index_2 | XXXXX  | 1      |
--------+--------+--------+
Index_3 | XXXXX  | 2      |
--------+--------+--------+
Index_4 | YYYYY  | 0      |
--------+--------+--------+
Index_5 | YYYYY  | 1      |
--------+--------+--------+
Index_6 | ZZZZZ  | 0      |
--------+--------+--------+
          [....]

我想要一个数据帧，如下所示：

          Col_A  | Col_B  
        +--------+------------------------------+
Index   
--------+--------+------------------------------+ 
Index_1 | XXXXX  | [Index_1, Index_2, Index_3]  | 
--------+--------+------------------------------+ 
Index_4 | YYYYY  | [Index_4, Index_5]           |
--------+--------+------------------------------+
Index_6 | ZZZZZ  | [Index_6, ...]               |
--------+--------+------------------------------+
          [....]

我的方法是按Col_A分组并迭代每组的行。或者更好地遍历数据框本身？性能很重要，因为数据帧非常大。

Answer 1

df.reset_index().groupby('Col_A') \
  .agg(dict(Index='first', Col_B=lambda x: list(x))) \
  .reset_index().set_index('Index')

Answer 2

无需迭代，您可以直接应用list：

df.groupby('Col_A')['Col_B'].apply(list)

请务必在系列中选择'Col_B'以应用list，否则list将返回数据框的列。

修改

您可以按如下方式存储原始（第一个）索引：

original_index = df.reset_index(drop=False).groupby('Col_A')['Index'].first()

然后将原始索引与之前的结果合并：

result = df.groupby('Col_A')['Col_B'].apply(list).join(original_index)

最后，您可以将索引与：

交换

result.reset_index(drop=False).set_index('Index')