Home > Software engineering >  Apply a function on a 2D array using a mask and a starting index
Apply a function on a 2D array using a mask and a starting index

Time:08-02

Let's say I have a 2D numpy array as follows:

x = array([
           [0, 7, 1, 6, 2, 3, 4],
           [9, 5, 1, 3, 5, 4, 8],
           [8, 5, 8, 1, 1, 2, 0],
           [5, 6, 3, 9, 8, 9, 1],
           [2, 9, 4, 6, 7, 6, 0]
          ])

I want to apply a function to that array starting at a specific index and using a mask.

For example, I want to apply the function value = randint(250, 255) starting at position [0,3] using the mask

 [
  [1, 1, 1],
  [1, 0, 1],
  [1, 0, 1],
  [1, 0, 1],
  [1, 1, 1]
 ]

which would give me

x = [
     [0, 7, 1, 254, 252, 253, 4],
     [9, 5, 1, 251, 5,   251, 8],
     [8, 5, 8, 255, 1,   252, 0],
     [5, 6, 3, 250, 8,   250, 1],
     [2, 9, 4, 252, 254, 255, 0]
    ]

P.S. Array x is randomly generated and can be [30, 15], [60, 30] or [120, 60]. The starting index will also be randomly selected

CodePudding user response:

A bit tricky, but you can extend the mask to full size, then use a flattened view to assign the changed values:

mask = np.array([[1, 1, 1],
                 [1, 0, 1],
                 [1, 0, 1],
                 [1, 0, 1],
                 [1, 1, 1]
                 ])

# extended mask
m = np.zeros_like(x, dtype=bool)
m[0:0 mask.shape[0], 3:3 mask.shape[1]] |= mask==1

# assign new values using a flattened view
x.ravel()[m.ravel()] = np.random.randint(250, 255, size=m.sum())

output:

array([[  0,   7,   1, 253, 250, 252,   4],
       [  9,   5,   1, 252,   5, 254,   8],
       [  8,   5,   8, 254,   1, 254,   0],
       [  5,   6,   3, 254,   8, 252,   1],
       [  2,   9,   4, 252, 254, 253,   0]])

intermediate m:

array([[False, False, False,  True,  True,  True, False],
       [False, False, False,  True, False,  True, False],
       [False, False, False,  True, False,  True, False],
       [False, False, False,  True, False,  True, False],
       [False, False, False,  True,  True,  True, False]])

CodePudding user response:

Why can't you just use a full mask?

from random import randint

x = [
  [0, 7, 1, 6, 2, 3, 4],
  [9, 5, 1, 3, 5, 4, 8],
  [8, 5, 8, 1, 1, 2, 0],
  [5, 6, 3, 9, 8, 9, 1],
  [2, 9, 4, 6, 7, 6, 0]
]

mask = [
  [0, 0, 0, 1, 1, 1, 0],
  [0, 0, 0, 1, 0, 1, 0],
  [0, 0, 0, 1, 0, 1, 0],
  [0, 0, 0, 1, 0, 1, 0],
  [0, 0, 0, 1, 1, 1, 0]
]

for i,row in enumerate(x):
  for j,_ in enumerate(row):
    if mask[i][j]:
      x[i][j] = randint(250,255)
print(x)
  • Related