mxnet.np.where

where(condition[, x, y])

Return elements chosen from x or y depending on condition.

Note

When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero(). The rest of this documentation covers only the case where all three arguments are provided.

Parameters

  • condition (ndarray) – Where True, yield x, otherwise yield y.

  • y (x,) – Values from which to choose. x, y and condition need to be broadcastable to some shape. x and y must have the same dtype.

Returns

out – An array with elements from x where condition is True, and elements from y elsewhere.

Return type

ndarray

Notes

If all the arrays are 1-D, where is equivalent to:

[xv if c else yv
for c, xv, yv in zip(condition, x, y)]

Examples

>>> a = np.arange(10)
>>> a
array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
>>> np.where(a < 5, a, 10*a)
array([ 0.,  1.,  2.,  3.,  4., 50., 60., 70., 80., 90.])

This can be used on multidimensional arrays too:

>>> cond = np.array([[True, False], [True, True]])
>>> x = np.array([[1, 2], [3, 4]])
>>> y = np.array([[9, 8], [7, 6]])
>>> np.where(cond, x, y)
array([[1., 8.],
       [3., 4.]])

The shapes of x, y, and the condition are broadcast together:

>>> x, y = onp.ogrid[:3, :4]
>>> x = np.array(x)
>>> y = np.array(y)
>>> np.where(x < y, x, 10 + y)  # both x and 10+y are broadcast
array([[10,  0,  0,  0],
       [10, 11,  1,  1],
       [10, 11, 12,  2]], dtype=int64)

>>> a = np.array([[0, 1, 2],
...               [0, 2, 4],
...               [0, 3, 6]])
>>> np.where(a < 4, a, -1)  # -1 is broadcast
array([[ 0.,  1.,  2.],
       [ 0.,  2., -1.],
       [ 0.,  3., -1.]])