Find the indices of non-zero elements and group by values

numpy, optimization, python

Solution

Here's an O(n log n) algorithm for your problem. The obvious looping solution is O(n), so for sufficiently large datasets this will be slower:

>>> 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])]

It will depend on the size of your data, but it is reasonably fast for largish data sets:

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

Problem

I wrote a code in python that takes a numpy matrix as input and returns a list of indices grouped by the corresponding values (i.e. output[3] returns all indices with value of 3). However, I lack the knowledge of writing vectorized code and had to do it using ndenumerate. This operation only took about 9 seconds which is too slow. The second idea that I had was using numpy.nonzero as follows: ``` for i in range(1, max_value): current_array = np.nonzero(input == i) # save in an array ``` This took 5.5 seconds and so it was a good improvement but still slow. Any way to do it without loops or optimized way to get the pairs of indices per value?

Original source