Generating multiple new arrays using numpy

Question

Say I have the following two different states:

state1 = [0, 1, 0, 1, 1]
state2 = [1, 1, 0, 0, 1]

And I want to generate n number of new states by only changing the values at where these states differ; and I want to do so randomly. So I did the following:

state1 = np.array(state1)
differ = np.where(state1 != state2)[0]  # --> array([0, 3], dtype=int64)

rand = [random.choices([1, 0], k=len(differ)) for _ in range(4)]
states = [state1.copy() for _ in range(4)]

for i in range(len(states)):
    states[i][differ] = rand[i]  # Will replace the values at locations where they are mismatched only

Out:

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

This does the job but I was wondering if there is a better way (to avoid 2 for loops), but this was the best I could come up with. Any ideas on how to do the same thing but better/cleaner?

Answer 1

You could use np.random.choice() to generate the states, while using a==b and a!=b as masks.

def new_states(a,b,n):
    return [
        a*(a==b) + np.random.choice(2,len(a))*(a!=b)
        for _ in range(n)
    ]

state1 = np.array([0, 1, 0, 1, 1])
state2 = np.array([1, 1, 0, 0, 1])
print(new_states(state1,state2,4))