假设我有一个 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 毫秒。尽管如此，我觉得它可以通过更好地利用矢量化（也许）来进一步改进。
有人可以指导我如何使其更简单和/或更快吗？任何帮助表示赞赏。谢谢！