gogurt gogurt - 2 months ago 12
Python Question

Drawing a random number from set of numbers excluding those given in certain sets

Suppose I have a two arrays:

import numpy as np
a = np.random.randint(0,10,10)
b = np.random.randint(0,10,10)


I want to generate another length-10 array whose i-th entry is a random integer drawn from the set (
{0...9}
minus the elements
a[i]
and
b[i]
).

Being a relative rookie when it comes to NumPy, I thought the easiest way to do this might be:


  1. obtain the set difference
    x = {0...9} - (a[i] union b[i])
    for each
    i

  2. do
    np.random.choice(x[i], 1)
    for each
    i



But I'm finding this a little tricky because I can't figure out how to map
setdiff1d
elementwise across 2 arrays. Is there an obvious way to do this in NumPy (i.e. ideally without having to resort to Python sets etc.)?

Answer

Here is one way:

In [87]: col = np.array((a, b)).T

In [88]: r = np.arange(10)

In [89]: np.ravel([np.random.choice(np.setdiff1d(r, i), 1) for i in col])
Out[89]: array([7, 8, 8, 6, 6, 8, 6, 5, 5, 6])

Or as a numpytonic approach:

In [101]: def func(x):                           
             return np.random.choice(np.setdiff1d(r, x), 1)
   .....: 

In [102]: np.apply_along_axis(func, 1, col).ravel()
Out[102]: array([6, 7, 9, 6, 4, 6, 7, 4, 0, 7])