Replace a single value with multiple values-CodePudding

Given a mask numpy array such as:

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

I want to replace each 1 with a target vector. Example:

target = np.array([5, 4, 3, 2, 1])
mask =   np.array([0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0,...])
output = np.array([0, 0, 5, 4, 3, 2, 1, 0, 5, 4, 3, 2, ...])

# Overlaps:
mask =   np.array([0, 0, 1, 0, 0, 0, 1, 0, 0, 0,...])
output = np.array([0, 0, 5, 4, 3, 2, 5, 4, 3, 2, ...])

Naivly, one can write this via the following (ignoring boundary problems):

output = np.zeros_like(mask)
for i, x in enumerate(mask):
  if x == 1:
    output[i:i len(target)] = target

I'm wondering, whether this is possible without resorting to a for loop?

CodePudding user response：

Not a full answer, but some thoughts: The for loop is O(n), where n = len(mask). We can use np.split to get that down to O(k), where k = number of 1s in mask:

def set_target(mask, target):
  output = []
  i, = np.where(mask == 1)
  for split in np.split(mask, i):
    if len(split) > len(target):
      split[:len(target)] = target
      output.append(split)
    else:
      output.append(target[:len(split)])
  return np.concatenate(output, 0)

CodePudding user response：

numpy supports assigning value for the same index multiple times in one go, like so:

mask =   np.array([0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0])
padding_idx = [2,3,4,5,6,5,6,7,8,9,8,9,10,11]
padding_values = [5,4,3,2,1,5,4,3,2,1,5,4,3,2]
mask[padding_idx] = padding_values
>>> mask
array([0, 0, 5, 4, 3, 5, 4, 3, 5, 4, 3, 2])

You just need to find out padding_idx and padding_values.

Note that padding_values = [5,4,3,2,1,5,4,3,2,1,5,4,3,2] has one value missing. So you need also to find a number of values missing. After that you can use broadcasting

vector = np.array([5,4,3,2,1])
N = len(vector)
mask =   np.array([0, 0, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0])
idx = np.flatnonzero(mask)
missing_values = len(mask) - idx[-1] - N

#Broadcast
padding_idx = np.flatnonzero(mask)[:,None]   np.arange(N)
padding_values = np.repeat(vector[np.newaxis, :], len(idx), axis=0)

#Flatten
padding_idx = padding_idx.ravel()[:missing_values]
padding_values = padding_values.ravel()[:missing_values]

#Go!
mask[padding_idx] = padding_values
>>> mask
array([0, 0, 5, 4, 3, 5, 4, 3, 5, 4, 3, 2])