根据另一个数组的值（按升序排序）将 NumPy 数组拆分为子数组

Question

假设我有两个 NumPy 数组

x = [[5, 2, 8],
     [4, 9, 1],
     [7, 8, 9],
     [1, 3, 5],
     [1, 2, 3],
     [1, 2, 4]]
y = [0, 0, 1, 1, 1, 2]

我想根据 y.

中的值有效地将数组 x 拆分为子数组

我想要的输出是

z_0 = [[5, 2, 8],
       [4, 9, 1]]
z_1 = [[7, 8, 9],
       [1, 3, 5],
       [1, 2, 3]]
z_2 = [[1, 2, 4]]

假设 y 从零开始并按升序排序，最有效的方法是什么？

注意：这个问题是这个问题的排序版本：

Answer 1

如果y是分组的（不用排序），可以用diff得到分割点：

indices = np.flatnonzero(np.diff(y)) + 1

您可以将它们直接传递给 np.split:

z = np.split(x, indices, axis=0)

如果你也想知道标签：

labels = y[np.r_[0, indices]]

Split a NumPy array into subarrays according to the values (sorted in ascending order) of another array