How to get the index of non-zero value in a numpy array quickly?

Question

Now I am writing a function which is to get the index of the non-zero values with following rules:

1. The expected result is a list. Each element denotes the index of a continue slice of non-zero values. So for a list of [0,0,0,1,1,1,0,1,1,0] , it should get the list [[3,4,5], [7,8]]
2. The index of different values in the list should be in separated list, that is, for the list of [0,0,1,1,1,2,2,1,1,0] , the expected result is [[2,3,4],[5,6],[7,8]] .

Do you have any idea? Thank you in advance!

Answer 1

With arr as input array and to have a list of arrays as output, you could do something like this -

# Store non-zero element indices
idx = np.where(arr)[0]

# Get indices where the shifts occur, i.e. positions where groups of identical 
# elements are separated. For this we perform differnetiation and look for 
# non-zero values and then get those positions. Finally, add 1 to compensate 
# for differentiation that would have decreased those shift indices by 1.
shift_idx = np.where(np.diff(arr[idx])!=0)[0]+1

# Split the non-zero indices at those shifts for final output
out = np.split(idx,shift_idx)

Sample input, output -

In [35]: arr
Out[35]: array([0, 0, 1, 1, 1, 2, 2, 1, 1, 0, 2, 2, 4, 3, 3, 3, 0])

In [36]: out
Out[36]: 
[array([2, 3, 4]),
 array([5, 6]),
 array([7, 8]),
 array([10, 11]),
 array([12]),
 array([13, 14, 15])]