Python - How to split an array based on the first column?

Question

I have below fake data. After reading it into array it will have shape (8, 3). Now I want to split the data based on the first column(ID) and return a list of array whose shape will be:[(3,3),(2,3),(3,3)]. I think np.split could do the job by assigning a 1-D array to "indices_or_sections" argument. But is there any more convenient way to do this?

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

You can achieve this by using a combination of np.split , sort , np.unique and 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]])]