在Python / numpy中，如何根据数组的第一列（即索引）解散/拆分数组？

Question

这个问题与如何根据numpy中的条件拆分数组密切相关？ 但我正在寻找一种更通用的方式来拆分给定未知数目的索引的数组：

import numpy as np
a=np.arange(10,40).reshape(10,3)
b=np.array([[1],[1],[1],[2],[2],[3],[3],[4],[4],[5]])
c=np.hstack((b,a))

array([[ 1, 10, 11, 12],
       [ 1, 13, 14, 15],
       [ 1, 16, 17, 18],
       [ 2, 19, 20, 21],
       [ 2, 22, 23, 24],
       [ 3, 25, 26, 27],
       [ 3, 28, 29, 30],
       [ 4, 31, 32, 33],
       [ 4, 34, 35, 36],
       [ 5, 37, 38, 39]])

我想根据第一列将其拆分为1x2x2数组，因此：

array([[[ 1, 10, 11, 12],
        [ 1, 13, 14, 15],
        [ 1, 16, 17, 18]],

       [[ 2, 19, 20, 21],
        [ 2, 22, 23, 24]],

       [[ 3, 25, 26, 27],
        [ 3, 28, 29, 30]],

       [[ 4, 31, 32, 33],
        [ 4, 34, 35, 36]],

       [[ 5, 37, 38, 39]]])

我是Python的新手，所以在此先感谢您的帮助！

Answer 1

在纯python中，您可以使用itertools.groupby来做到这一点：

>>> from operator import itemgetter
>>> from itertools import groupby
>>> from pprint import pprint
>>> pprint ([list(g) for k, g in groupby(c, key=itemgetter(0))])
[[array([ 1, 10, 11, 12]), array([ 1, 13, 14, 15]), array([ 1, 16, 17, 18])],
 [array([ 2, 19, 20, 21]), array([ 2, 22, 23, 24])],
 [array([ 3, 25, 26, 27]), array([ 3, 28, 29, 30])],
 [array([ 4, 31, 32, 33]), array([ 4, 34, 35, 36])],
 [array([ 5, 37, 38, 39])]]

使用NumPy ：

>>> e, inds = np.unique(c[:,0], return_index=True)
>>> np.split(c, inds)[1:]
[array([[ 1, 10, 11, 12],
       [ 1, 13, 14, 15],
       [ 1, 16, 17, 18]]),
 array([[ 2, 19, 20, 21],
       [ 2, 22, 23, 24]]),
 array([[ 3, 25, 26, 27],
       [ 3, 28, 29, 30]]),
 array([[ 4, 31, 32, 33],
       [ 4, 34, 35, 36]]),
 array([[ 5, 37, 38, 39]])]