python - 在 Python/numpy 中，如何根据数组的第一列(索引)解开/拆分数组？

Question

这个问题与How to split an array according to a condition in 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]])]