How do I get a list of indices of non zero element

问题:

I have a list that will always contain only ones and zeroes. I need to get a list of the non-zero indices of the list:

a = [0, 1, 0, 1, 0, 0, 0, 0]
b = []
for i in range(len(a)):
    if a[i] == 1:  b.append(i)
print b

What would be the 'pythonic' way of achieving this ?

回答1:

[i for i, e in enumerate(a) if e != 0]

回答2:

Not really a "new" answer but numpy has this built in as well.

import numpy as np
a = [0, 1, 0, 1, 0, 0, 0, 0]
nonzeroind = np.nonzero(a)[0] # the return is a little funny so I use the [0]
print nonzeroind
[1 3]

回答3:

Since THC4k mentioned compress (available in python2.7+)

>>> from itertools import compress, count
>>> x = [0, 1, 0, 1, 0, 0, 0, 0]
>>> compress(count(), x)
<itertools.compress object at 0x8c3666c>   
>>> list(_)
[1, 3]

回答4:

Just wished to add explanation for 'funny' output from the previous asnwer. Result is a tuple that contains vectors of indexes for each dimension of the matrix. In this case user is processing what is considered a vector in numpy, so output is tuple with one element.

import numpy as np
a = [0, 1, 0, 1, 0, 0, 0, 0]
nonzeroind = np.nonzero(a) 
print nonzeroind
(array([1, 3]),)