创建字典的最快方法是在二维 numpy 数组中按分组样式查找？

Question

假设我有一个 2D numpy 数组，其值对应于一个标签或类。例如，如果 A = [[0, 0, 1, 1], [1, 1, 1, 0]，则位置(0, 0), (0, 1), (1, 3)对应于“0”(0, 2), (0, 3), (1, 0), etc类并对应于“1”类。这是一个非常简单的例子，但总的来说，我会处理包含更多项的矩阵。

我想要做的本质上是构建一个字典，其中一个键对应于每个类，其对应的值是一个元组列表，其中每个元组对应于输入矩阵的一个位置，其值为键。换句话说，按输入矩阵的值对输入矩阵进行分组，并获得每个唯一值出现的位置列表。

现在，我有以下代码：

S = {i: [] for i in range(A.max() + 1)}
for i in range(A.shape[0]):
    index = np.arange(A[i].shape[0])
    sort_idx = np.argsort(A[i])
    cnt = np.bincount(A[i])
    result = np.split(index[sort_idx], np.cumsum(cnt[:-1]))
    for j, k in enumerate(result):
        S[j] += [(i, z) for z in k]

其中 A 是我的输入矩阵。在 500x500 矩阵上运行平均需要大约 0.4 毫秒。尽管如此，我觉得它可以通过更好地利用矢量化（也许）来进一步改进。

有人可以指导我如何使其更简单和/或更快吗？任何帮助表示赞赏。谢谢！

Answer 1

您可以更简单地使用np.argwhereand做到这一点np.unique：

S = {}
for key in np.unique(A):
    S[key] = np.argwhere(A==key)

请注意，这将返回一个 2D numpy 数组。