Pandas：使用组索引创建数组

Question

我正在寻找一种高效的、全熊猫的方式来创建具有组号的数组（对于原始数据框中的每一行，我想要一个数字来告诉我该行属于哪个组）：

df = pandas.DataFrame({'a': [1, 1, 1, 2, 2, 2], 'b': [1, 2, 1, 1, 2, 1]})
groups = df.groupby(['a', 'b'])
group_names = sorted(groups.groups.keys())
group_indices = np.array(df.index)
for index, group in enumerate(group_names):
    group_indices[groups.indices[group]] = index

在哪里

In : df 
Out]:
   a  b
0  1  1
1  1  2
2  1  1
3  2  1
4  2  2
5  2  1    

In : groups.indices
Out:
{(1, 1): array([0, 2]),
 (1, 2): array([1]),
 (2, 1): array([3, 5]),
 (2, 2): array([4])}

In : group_indices
Out: array([0, 1, 0, 2, 3, 2])

我的问题是，如果 df 大约为 20000x100（64 位浮点数）并且我按两列分组，我的内存使用量将超过 6 GB。这超出了我的预期。

Answer 1

索引已经嵌入到 groupby 对象中

In [52]: groups.grouper.levels
Out[52]: [Int64Index([1, 2], dtype=int64), Int64Index([1, 2], dtype=int64)]

In [53]: groups.grouper.labels
Out[53]: [array([0, 0, 0, 1, 1, 1]), array([0, 1, 0, 0, 1, 0])]

In [57]: l = groups.grouper.labels

In [58]: zip(*l)
Out[58]: [(0, 0), (0, 1), (0, 0), (1, 0), (1, 1), (1, 0)]

In [18]: groups.grouper.group_info
Out[18]: (array([0, 1, 0, 2, 3, 2]), array([0, 1, 2, 3]), 4)

简单的查找操作，因为这些已经在分组对象上计算了

In [19]: groups.grouper.group_info[0]
Out[19]: array([0, 1, 0, 2, 3, 2])