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

Question

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

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

我具有与CCD_ 3的行相对应的另一个标签阵列CCD_。例如：

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

我想根据A的标签将其拆分为3个阵列，因此结果将是：

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

有没有内置的numpy函数可以实现这一点？

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

Answer 1

这里有一种方法：

1. hstack二者共同组成阵列
2. sortarray由the last column
3. split基于unique值index的array

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

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

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 3

您也可以使用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 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]])]