python - 创建在 2D 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)等对应类'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.argwhere 和 np.unique 更简单地完成此操作:

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

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