I need to fill a numpy array of three elements with random integers such that the sum total of the array is three (e.g. [0,1,2]).
By my reckoning there are 10 possible arrays:
111, 012, 021, 102, 120, 201, 210, 300, 030, 003
My ideas is to randomly generate an integer between 1 and 10 using randint, and then use a look-up table to fill the array from the above list of combinations.
Does anyone know of a better approach?
Here is how I did it:
>>> import numpy as np
>>> a=np.array([[1,1,1],[0,1,2],[0,2,1],[1,0,2],[1,2,0],[2,0,1],[2,1,0],[3,0,0],[0,3,0],[0,0,3]])
>>> a[np.random.randint(0,10)]
array([1, 2, 0])
>>> a[np.random.randint(0,10)]
array([0, 1, 2])
>>> a[np.random.randint(0,10)]
array([1, 0, 2])
>>> a[np.random.randint(0,10)]
array([3, 0, 0])
Here’s a naive programmatic way to do this for arbitrary array sizes/sums:
def n_ints_summing_to_v(n, v):
elements = (np.arange(n) == np.random.randint(0, n)) for i in range(v))
return np.sum(elements, axis=0)
This will, of course, slow down proportionally to the desired sum, but would be ok for small values.
Alternatively, we can phrase this in terms of drawing samples from the Multinomial distribution, for which there is a function available in NumPy (see here), as follows:
def n_ints_summing_to_v(n, v):
return np.random.multinomial(v, ones((n)) / float(n))
This is a lot quicker!
This problem can be solved in the generic case, where the number of elements and their sum are both configurable. One advantage of the solution below is that it does not require generating a list of all the possibilities. The idea is to pick random numbers sequentially, each of which is less than the required sum. The required sum is reduced every time you pick a number:
import numpy
def gen(numel = 3, sum = 3):
arr = numpy.zeros((numel,), dtype = numpy.int)
for i in range(len(arr) - 1): # last element must be free to fill in the sum
arr[i] = numpy.random.randint(0, sum + 1)
sum -= arr[i]
if sum == 0: break # Nothing left to do
arr[-1] = sum # ensure that everything adds up
return arr
print(gen())
This solution does not guarantee that the possibilities will all occur with the same frequency. Among the ten possibilities you list, four start with 0, three with 1, two with 2 and one with 3. This is clearly not the uniform distribution that numpy.random.randint() provides for the first digit.
