python - Creating an Array with an equal number of 0 and 1 in a random order?

Question

Before I explain, here is my code for reference:

import numpy as np 

arrayteam = [[3,3,3,3,3,3],[2,2,2,2,2,2],[1,1,1,1,1,1]]
#nteams = 3
#nsubteam = 2

#newnums = np.zeros((len(arrayteam),len(arrayteam[0])))

subteam = [] 
for i in range(len(arrayteam)):
    subteam.append(np.random.choice([0,1],size=len(arrayteam[i]),p=[0.5,0.5]))

print (subteam)

Here is the output:

[array([0, 1, 0, 1, 1, 1]), array([1, 1, 0, 1, 0, 0]), array([1, 0, 1, 1, 1, 1])]

As you can see, it randomly chooses 0s and 1s which is what I want, however the number of 0s and 1s is unequal in each array, obviously because I have it as p=0.5, so there's a 50% chance it will choose 0 or 1. I want to have it so that there are 3 zeros and 3 ones in each array, but they occur in a random order. How can I do this?

Also, how can I end up changing exactly where I want the zeros and ones to occur? For example, what if I want the first 3 numbers in the array to be 0s, and the second 3 to be ones? Or what if I want them to alternate?

Answer 1

Use random.shuffle. Note that it shuffles in-place, and returns None:

import random
length = 3
subteam = [0] * length + [1] * length
random.shuffle(subteam)
print(subteam)