首页 > 解决方案 > 熊猫如何在排序的 groupby 中找到组 n?

问题描述

我有一个巨大的 Pandas 数据框转换为用于深度学习的张量。现在我想访问我的数据框和数组中的组 n。

例如,如何使用给定的一组排序键访问最终成为数组中第 3 组的数据帧组?

在大型数据集上,将数据帧转换为数组非常慢,因为在数百万行中大约有 200k 个组。所以这只能完成一次,而不是即时完成,而且我没有组列表的泡菜。

import pandas as pd
import numpy as np

np.random.seed(0)

df = pd.DataFrame({"a"     : np.random.normal(0, 1, 100),
                   "index1": np.random.randint(0, 5, 100),
                   "index2": np.random.randint(0, 5, 100)})

grouped_df = df.groupby(["index1", "index2"])

# convert dataframe to array for e.g. deep learning but never do this operation again (too slow)
array = np.array([group["a"].values for _, group in grouped_df])

# fetch the same sample from the array and the df
array_n = array[3] # this is trivial

# how can I do this in my df?
# grouped_df[3] obviously doesn't work.

标签: pythonpandas

解决方案


您可以使用属性获取组键列表grouped_df.groups,该属性返回一个字典。然后使用以下函数获取该字典的键keys()

In [27]: grouped_df.groups.keys()
Out[27]: dict_keys([(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (2, 0), (2, 1), (2, 3), (2, 4), (3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (4, 0), (4, 1), (4, 2), (4, 3), (4, 4)])

(这些键的值(index1, index2)对应于该组的 index1/index2 值。)

您用于填充的列表推导array将按顺序遍历这些组键,因此结果数组的元素 4 对应于第四个键:

In [28]: list(grouped_df.groups.keys())[3]
Out[28]: (0, 3)

(表示index1=0index2=3)。现在您可以将该键分配给一个变量,并使用以下方法获取与该键对应的组grouped_df.get_group()

In [29]: my_key = list(grouped_df.groups.keys())[3]

In [30]: grouped_df.get_group(my_key)
Out[30]:
           a  index1  index2
0   1.764052       0       3
14  0.443863       0       3
56  0.066517       0       3
58 -0.634322       0       3
65 -0.401781       0       3
69  0.051945       0       3

最后,从结果中抓取列“a”,并以与列表推导相同的方式获取值:

In [31]: grouped_df.get_group(my_key)['a'].values
Out[31]:
array([ 1.76405235,  0.44386323,  0.06651722, -0.63432209, -0.40178094,
        0.0519454 ])

或者,作为一个有点凌乱的单线,

In [32]: grouped_df.get_group(list(grouped_df.groups.keys())[3])['a'].values
Out[32]:
array([ 1.76405235,  0.44386323,  0.06651722, -0.63432209, -0.40178094,
        0.0519454 ])

推荐阅读