In [13]: arr = np.array([[6, 1, 4],
...: [4, 8, 4],
...: [6, 3, 5]])
In [14]: idx = np.nonzero(arr>4)
In [15]: idx
Out[15]: (array([0, 1, 2, 2]), array([0, 1, 0, 2]))
This tuple can be used directly to index the array, the result being a 1d array of the >4 values:
In [16]: arr[idx]
Out[16]: array([6, 8, 6, 5])
it can also be used to modify values
In [17]: arr1 = arr.copy()
In [18]: arr1[idx] = 10
In [19]: arr1
Out[19]:
array([[10, 1, 4],
[ 4, 10, 4],
[10, 3, 10]])
You can get an array index pairs with:
In [20]: np.argwhere(arr>4)
Out[20]:
array([[0, 0],
[1, 1],
[2, 0],
[2, 2]])
argwhere
just applies transpose
to the where
:
In [21]: np.transpose(idx)
Out[21]:
array([[0, 0],
[1, 1],
[2, 0],
[2, 2]])
arr[idx]
is the equivalent of:
In [22]: arr[idx[0],idx[1]]
Out[22]: array([6, 8, 6, 5])
np.take
is a way of indexing along just one axis at at time, so isn't very useful in this context. The indices in idx
are meant to be used together, not individually.