在 numpy 中,如何在不使用 for 循环的情况下有效地构建从每个唯一值到其索引的映射

In numpy, how to efficiently build a mapping from each unique value to its indices, without using a for loop

在 numpy 中,如何在不使用 for 循环的情况下高效地构建从每个唯一值到其索引的映射

我考虑了以下备选方案,但它们对于我的用例来说不够有效,因为我使用大型数组。

第一种选择,需要使用 for 循环遍历数组,这在考虑大型 numpy 数组时可能会很慢。

import numpy as np
from collections import defaultdict
a = np.array([1, 2, 6, 4, 2, 3, 2])
inv = defaultdict(list)
for i, x in enumerate(a):
    inv[x].append(i)

第二种方法效率不高,因为它需要多次遍历数组:

import numpy as np
a = np.array([1, 2, 6, 4, 2, 3, 2])
inv = {}
for x in np.unique(a):
    inv[x] = np.flatnonzero(a == x)

编辑:我的 numpy 数组由整数组成,用于图像分割。我也在找skimage的方法,没找到。

我建议您查看 numba,它可以显着加快 python 上的 numpy 代码 - 它支持 numpy.invert()numpy.unique() - documentation

这是一个很好的视频,解释了如何使用来自 youtube 的 numba - Make Python code 1000x Faster with Numba

这应该有效。

a = np.array((1, 2, 6, 2, 4, 7, 25, 6))
fwd = np.argsort(a)
asorted = a[fwd]
keys = np.unique(asorted)
lower = np.searchsorted(asorted, keys)
# or higher = itertools.chain(lower[1:], (len(asorted),))
higher = np.append(lower[1:], len(asorted))
inv = {key: fwd[lower_i:higher_i]
       for key, lower_i, higher_i
       in zip(keys, lower, higher)}

assert all(np.all(a[indices] == key)
           for key, indices in inv.items())

它的运行时间类似于 O(n log(n))。唯一剩下的循环是构建字典的循环。当然,该步骤是可选的。

从纯算法的角度来看,您的第一种方法 (defaultdict(list)) 会更好,因为它以聚合线性时间运行,但当然 python 开销可能很大。