python遍历指定的数组元素

Question

I am using numpy.ndenumerate for array element iteration:

a = np.ones((10,10,10))
for (x,y,z), value in numpy.ndenumerate(a):
    do stuff with x,y,z and value

How can I iterate only over specified element values? For example if the array is binary, how do I iterate over all ones?

Currently I am using something like:

specified = np.nonzero(a) #or np.where(a == some_value)
for i in xrange(np.count_nonzero(a)):
    do stuff with x,y,z = specified[0][i], specified[1][i], specified[2][i] 
    and value = a[x,y,z]

This is obviously pretty ugly and I am shure there is a nice short way.

Answer 1

I don't see anything wrong with np.nonzero(a) . However, what's returned from np.nonzero is a tuple of elements where each element in the tuple is the indices that are non-zero for a particular dimension. Specifically, the first element of the tuple are the row locations of what is non-zero, the second element is the column locations of what is non-zero, the third element is the slice locations of what is non-zero, etc. Therefore, for each grouping of elements in the same index for each tuple element, you get a group of coordinates that tell you which locations are non-zero in your array.

Once you obtain these, you can simply zip over them and extract your values:

specified = np.nonzero(a)
for (x,y,z) in zip(*specified):
     # do your stuff with a[x,y,z] and x,y,z...
     value = a[x,y,z]
     #....
     #....

The above assumes that specified has three tuple elements corresponding to a being three dimensional. At each iteration of the loop, we grab unique (x,y,z) triplets that correspond to non-zero entries in a . You can go ahead and do your processing on each triplet.

---

However, the above code assumes that you're going to do something with both the (x,y,z) triplets and the values themselves within the same for loop body. If it is your intention to extract out the values only, then Boolean indexing is the better way to go. Specifically, you can just do something like this:

values = a == some_value

values will contain a Boolean array of the same size as a that indicates where a value is equal or not equal to some_value . From there, you can perhaps do something like this:

a[values] = some_other_value

This replaces all values in a that equal to some_value with some_other_value . If it's your intention to simply access values and do some manipulation on these values, then the above is the better way to go. However, I can't say that for sure as I don't know what "stuff" you're intending to do inside your for loop body.