Python - 如何根据第一列拆分数组？

Question

我有以下虚假数据。 将其读入数组后，它将具有形状 (8, 3)。 现在我想根据第一列（ID）拆分数据并返回一个数组列表，其形状为：[（3,3），（2,3），（3,3）]。 我认为 np.split 可以通过将一维数组分配给“indices_or_sections”参数来完成这项工作。 但是有没有更方便的方法来做到这一点？

1   700 35
1   700 35
1   700 35
2   680 25
2   680 25
3   750 40
3   750 40
3   750 40

Answer 1

您可以通过使用np.split 、 sort 、 np.unique和np.cumsum的组合来实现此目的。

>>> a = [[1, 700, 35],
...      [1, 700, 35],
...      [1, 700, 35],
...      [2, 680, 25],
...      [2, 680, 25],
...      [3, 750, 40],
...      [3, 750, 40],
...      [3, 750, 40]]
>>> a = np.array(a)
>>> # sort the array by first column. 
>>> a = a[a[:,0].argsort()]
>>> np.split(a, np.cumsum(np.unique(a[:, 0], return_counts=True)[1])[:-1])
[array([[  1, 700,  35],
       [  1, 700,  35],
       [  1, 700,  35]]), array([[  2, 680,  25],
       [  2, 680,  25]]), array([[  3, 750,  40],
       [  3, 750,  40],
       [  3, 750,  40]])]