Is there built in function in numpy to iterate advanced in 3d array-CodePudding

I wanna make a function that takes an array as its first parameter takes an arbitrary sized and shaped arr array and overwrites all its values that are in the given [a,b] interval to be equal to c. The a, b, c numbers are given to the function as parameters.like input and output below

arr = np.array([[[5., 2., -5.], [4., 3., 1.]]])
overwrite_interval(arr, -2., 2., 100.) -> ndarray([[[5., 100., -5.], [4., 3., 100.]]])

def overwrite_interval(arr , a , b , c):  
  for i in arr[:,:]: 
  arr[a,b] = c

arr = np.array([[[5., 2., -5.], [4., 3., 1.]]])
assert overwrite_interval(arr, -2., 2., 100.) #-> ndarray([[[5., 100., -5.], [4., 3., 100.]]])

CodePudding user response：

I think the way you've worded your question doesn't line up with the example you've given. Firstly, the example array you've given is 3D, not 2D. You can do

>>> arr.shape
(1,2,3)
>>> arr.ndim
3

Presumably this is a mistake, and you want your array to be 2D, so you would do

arr = np.array([[5., 2., -5.], [4., 3., 1.]])

instead.

Secondly, if a and b are values that, if an element is between then to set that element to value c rather than a and b being indexes, then the np.where function is great for this.

def overwrite_interval(arr , a , b , c):
    inds = np.where((arr >= a) * (arr <= b))
    arr[inds] = c
    return arr

np.where returns a tuple, so sometimes it can be easier to work with boolean arrays directly. In which case, the function would look like this

def overwrite_interval(arr , a , b , c):
    inds = (arr >= a) * (arr <= b)
    arr[inds] = c
    return arr

Does this work for you, and is this your intended meaning? Note that the solution I've provided would work as is if you still meant for the initial array to be a 3D array.