Python Random Array of 0s and 1s [duplicate]

Question

I want to randomly produce an array of n ones and m zeros.

I thought of this solution:

1. produce the ones array (np.ones)
2. produce the zeros array (np.zeros)
3. combine them to one array (np.hstack)
4. shuffle the resulting array (np.random.shuffle)

Seems to be not natural as a solution. Some pythonic ideas?

Answer 1

Your solution seems reasonable. It states exactly what it's doing, and does it clearly.

Let's compare your implementation:

a = np.hstack((np.ones(n), np.zeros(m)))
np.random.shuffle(a)

… with an obvious alternative:

a = np.ones(n+m)
a[:m] = 0
np.random.shuffle(a)

That might save a bit of time not allocating and moving hunks of data around, but it takes a bit more thought to understand.

And doing it in Python instead of in NumPy:

a = np.array([1]*n + [0]*m)
np.random.shuffle(a)

… might be a little more concise, but it seems less idiomatically NumPy (in the same way that np.array([1]*n) is less idiomatic than np.ones(n)), and it's going to be slower and use more memory for no good reason. (You could improve the memory by using np.fromiter, but then it's pretty clearly not going to be more concise.)

Of course if you're doing this more than once, the real answer is to factor it out into a function. Then the function's name will explain what it does, and almost any solution that isn't too tortured will be pretty easy to understand…

Answer 2

I'd make an array of n ones and m zeros as

a = np.array([1] * n + [0] * m)

Then I'd call np.random.shuffle() on it.