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]])]