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

Question

我有一个 2D 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 内置函数来实现这一点？

目前，我的解决方案相当难看，包括在 for-loop 中重复调用 numpy.where，并将索引元组切片以仅包含行。

Answer 1

这是一种方法：

1. hstack 两个数组在一起。
2. sortarraythe last column
3. split array 基于 unique 值 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]])]

Answer 2

您也可以为此使用 Pandas，因为它专为标记数据而设计，并且具有强大的 groupby 方法。

import pandas as pd
index = pd.Index(B, name='label')
df = pd.DataFrame(A, index=index)
groups = {k: v.values for k, v in df.groupby('label')}
print(groups)

这会生成一个包含分组值数组的字典：

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

对于数组列表，您可以这样做：

groups = [v.values for k, v in df.groupby('label')]

Answer 3

如果输出总是可以是一个数组，因为标签是均匀分布的，你只需要按标签对数据进行排序：

idx = B.argsort()
n = np.flatnonzero(np.diff(idx))[0] + 1
result = A[idx].reshape(n, A.shape[0] // n, A.shape[1])

如果标签分布不均，则必须在外部维度中列出：

_, indices, counts = np.unique(B, return_counts=True, return_inverse=True)
result = np.split(A[indices.argsort()], counts.cumsum()[:-1])

使用np.where 的等效项不是很有效，但您可以不使用循环：

b, idx = np.unique(B, return_inverse=True)
mask = idx[:, None] == np.arange(b.size)
result = np.split(A[idx.argsort()], np.count_nonzero(mask, axis=0).cumsum()[:-1])

您可以同时计算所有标签的掩码，并通过计算每个类别中匹配元素的数量 (np.count_nonzero(mask, axis=0).cumsum()) 将其应用于已排序的 A (A[idx.argsort()])。因为np.split 总是添加一个隐式总索引，所以最后一个索引被从累积总和中去除。

Answer 4

这可能是最简单的方法：

groups = [A[B == label, :] for label in np.unique(B)]
print(groups)

输出：

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