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), 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 毫秒。尽管如此，我觉得可以通过更好地利用矢量化（也许）进一步改进它。

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

标签： pythonarraysnumpygroupingvectorization

Answer 1

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

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

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