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