Fairly distribute predefined value using binary array as mask

Question

While answering this question, I was wondering what is the most efficient/simple way to distribute a predefined value across the 1 values of a binary array. The value should be distributed as fairly as possible (the remainder is distributed randomly so that each cell can eventually get one extra).

For instance, distributing 4 in [1, 0, 1] would result in [2, 0, 2], while distributing 5 could give either [3, 0, 2] or [2, 0, 3]. The distribution should work on an array of arbitrary dimensions.

Let's call distribute the function. It should take as parameters value the value to be distributed and array the binary array and return an array of same shape with the distributed value.

def distribute(value, array):
    ...
    return # array of same shape with fairly distributed value

>>> distribute(10, np.array([0,1]))
array([ 0, 10])

>>> distribute(10, np.array([1,0,1]))
array([5, 0, 5])

>>> distribute(9, np.array([1,0,1]))
array([4, 0, 5]) # could also be array([5, 0, 4])

>>> distribute(9, np.array([[1,0],[0,1]]))
array([[5, 0],
       [0, 4]])

>>> distribute(10, np.diag(np.ones(4)).reshape(2,2,-1))
array([[[3., 0., 0., 0.],
        [0., 2., 0., 0.]],

       [[0., 0., 3., 0.],
        [0., 0., 0., 2.]]])

See my proposed solution in the answers. Please propose alternative approaches, the goal is to find the most efficient and/or simple.

Answer 1

def distribute(value, array):
    shape = array.shape # save the original shape of the array
    div, mod = map(int, divmod(value, np.sum(array))) # calculate equal split and remainder, ensuring int type
    array = (array*div).flatten() # distribute evenly and flatten
    array[np.random.choice(np.flatnonzero(array), mod, replace=False)] += 1 # allocate remainder randomly
    return array.reshape(shape) # return the array reshaped as original