查找非零元素的索引并按值分组

这是针对您的问题的O（n log n）算法。显而易见的循环解决方案是O（n），因此对于足够大的数据集，这会更慢：

>>> a = np.random.randint(3, size=10)>>> aarray([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