我用python编写了一个代码，该代码以numpy矩阵作为输入并返回按相应值分组的索引列表（即output [3]返回值为3的所有索引）。但是，我缺乏编写矢量化代码的知识，因此不得不使用ndenumerate来完成。此操作仅花费了大约9秒钟，这太慢了。
我的第二个想法是使用numpy.nonzero，如下所示：
for i in range(1, max_value):
current_array = np.nonzero(input == i)
# save in an array
这花费了5.5秒，因此这是一个不错的改进，但仍然很慢。有没有循环的优化方法或获取每个值的索引对的优化方法？
这是针对您的问题的O（n log n）算法。显而易见的循环解决方案是O（n），因此对于足够大的数据集，这会更慢：
>>> a = np.random.randint(3, size=10)
>>> a
array([1, 2, 2, 0, 1, 0, 2, 2, 1, 1])
>>> index = np.arange(len(a))
>>> sort_idx = np.argsort(a)
>>> cnt = np.bincount(a)
>>> np.split(index[sort_idx], np.cumsum(cnt[:-1]))
[array([3, 5]), array([0, 4, 8, 9]), array([1, 2, 6, 7])]
这将取决于您的数据大小，但是对于大型数据集来说，这是相当快的：
In [1]: a = np.random.randint(1000, size=1e6)
In [2]: %%timeit
...: indices = np.arange(len(a))
...: sort_idx = np.argsort(a)
...: cnt = np.bincount(a)
...: np.split(indices[sort_idx], np.cumsum(cnt[:-1]))
...:
10 loops, best of 3: 140 ms per loop
本文收集自互联网，转载请注明来源。
如有侵权，请联系[email protected] 删除。
我来说两句