python - 基于单独的标签数组拆分 numpy 2D 数组

Question

我有一个二维 numpy 数组 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 -loop，并将索引元组切片以仅包含行。

Answer 1

这是一种方法:

hstack两者的阵列在一起。

sort array来自 the last column

split array基于 unique值(value) index

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