基于单独的标签数组分割numpy 2D数组

Question

我有一个 2D numpy array A。例如：

A = np.array([[1, 2],
              [3, 4],
              [5, 6],
              [7, 8],
              [9, 0]])

我有另一个与B的行相对应的标签数组A。例如：

B = np.array([0, 1, 2, 0, 1])

我想A根据标签分成 3 个数组，所以结果是：

[[[1, 2],
  [7, 8]],
 [[3, 4],
  [9, 0]],
 [[5, 6]]]

是否有任何 numpy 内置函数可以实现此目的？

现在，我的解决方案相当丑陋，涉及numpy.where在for循环中重复调用，并对索引元组进行切片以仅包含行。

Answer 1

这是一种方法：

1. hstack两个数组在一起。
2. sort由array​the last column
3. split基于价值array​uniqueindex

a = np.hstack((A,B[:,None]))
a = a[a[:, -1].argsort()]
a = np.split(a[:,:-1], np.unique(a[:, -1], return_index=True)[1][1:])

输出：

[array([[1, 2],
        [7, 8]]),
 array([[3, 4],
        [9, 0]]),
 array([[5, 6]])]