Как разделить массив на основе первого столбца?

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