使用Numpy / Python将概率函数分别应用于数组的特定元素

Question

I'm fairly new to Python/Numpy. What I have here is a standard array and I have a function which I have vectorized appropriately.

def f(i):
     return np.random.choice(2,1,p=[0.7,0.3])*9
f = np.vectorize(f)

Defining an example array:

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

With the vectorized function, f, I would like to evaluate f on each cell on the array with a value of 0.

I am trying to leave for loops as a last resort. My arrays will eventually be larger than 100 by 100, so running each cell individually to look and evaluate f might take too long.

I have tried:

print f(array[array==0])

Unfortunately, this gives me a row array consisting of 5 elements (the zeroes in my original array).

Alternatively I have tried,

array[array==0] = f(1)

But as expected, this just turns every single zero element of array into 0's or 9's.

What I'm looking for is somehow to give me my original array with the zero elements replaced individually. Ideally, 30% of my original zero elements will become 9 and the array structure is conserved.

Thanks

Answer 1

The reason your first try doesn't work is because the vectorized function handle, let's call it f_v to distinguish it from the original f , is performing the operation for exactly 5 elements: the 5 elements that are returned by the boolean indexing operation array[array==0] . That returns 5 values, it doesn't set those 5 items to the returned values. Your analysis of why the 2nd form fails is spot-on.

If you wanted to solve it you could combine your second approach with adding the size option to np.random.choice :

array = np.array([[1,1,0],[0,1,0],[0,0,1]])
mask = array==0
array[mask] = np.random.choice([18,9], size=mask.sum(), p=[0.7, 0.3])
# example output:
# array([[ 1,  1,  9],
#        [18,  1,  9],
#        [ 9, 18,  1]])

There was no need for np.vectorize : the size option takes care of that already.