Why not simply use masking here:
z[z % 3 == 0]
For your sample matrix, this will generate:
>>> z[z % 3 == 0]
array([0, 3, 6])
If you pass a matrix with the same dimensions with booleans as indices, you get an array with the elements of that matrix where the boolean matrix is True.
This will furthermore work more efficient, since you do the filtering at the numpy level (whereas list comprehension works at the Python interpreter level).
Source for argwhere
def argwhere(a):
"""
Find the indices of array elements that are non-zero, grouped by element.
...
"""
return transpose(nonzero(a))
np.where is the same as np.nonzero.
In [902]: z=np.arange(9).reshape(3,3)
In [903]: z%3==0
Out[903]:
array([[ True, False, False],
[ True, False, False],
[ True, False, False]], dtype=bool)
In [904]: np.nonzero(z%3==0)
Out[904]: (array([0, 1, 2], dtype=int32), array([0, 0, 0], dtype=int32))
In [905]: np.transpose(np.nonzero(z%3==0))
Out[905]:
array([[0, 0],
[1, 0],
[2, 0]], dtype=int32)
In [906]: z[[0,1,2], [0,0,0]]
Out[906]: array([0, 3, 6])
z[np.nonzero(z%3==0)] is equivalent to using I,J as indexing arrays:
In [907]: I,J =np.nonzero(z%3==0)
In [908]: I
Out[908]: array([0, 1, 2], dtype=int32)
In [909]: J
Out[909]: array([0, 0, 0], dtype=int32)
In [910]: z[I,J]
Out[910]: array([0, 3, 6])
In [31]: a= np.array([1,2,3,4,10,12,13])
In [32]: a
Out[32]: array([ 1, 2, 3, 4, 10, 12, 13])
The 2 conditions individually:
In [33]: a<4
Out[33]: array([ True, True, True, False, False, False, False])
In [34]: a>10
Out[34]: array([False, False, False, False, False, True, True])
Join them with logical or:
In [35]: (a<4) | (a>10)
Out[35]: array([ True, True, True, False, False, True, True])
where is a tuple (size 1 because there's 1 dimension):
In [36]: np.where((a<4) | (a>10))
Out[36]: (array([0, 1, 2, 5, 6]),)
That tuple can be used directly to index a:
In [37]: a[_]
Out[37]: array([ 1, 2, 3, 12, 13])
argwhere gives the same indices but in a 'vertical' format:
In [38]: np.argwhere((a<4) | (a>10))
Out[38]:
array([[0],
[1],
[2],
[5],
[6]])
It may be easier to visualize, but isn't so useful for indexing (unless you do it iteratively).
With the logical and, the result is an empty index - no value meets both conditions:
In [39]: np.where((a<4) & (a>10))
Out[39]: (array([], dtype=int64),)
You can use np.where :
c = np.where((a < 4) & (a > 10),a, a)
c
array([ 1, 2, 3, 4, 10, 12, 13])