array rows where the random-integer elements may have different ranges

Question

Consider the following code fragment:

import numpy as np

mask = np.array([True, True, False, True, True, False])
val = np.array([9, 3])
arr = np.random.randint(1, 9, size = (5,len(mask)))

As expected, we get an array of random integers, 1 to 9, with 5 rows and 6 columns as below. The val array has not been used yet.

[[2, 7, 6, 9, 7, 5],
 [7, 2, 9, 7, 8, 3],
 [9, 1, 3, 5, 7, 3],
 [5, 7, 4, 4, 5, 2],
 [7, 7, 9, 6, 9, 8]]

Now I'll introduce val = [9, 3].
Where mask = True, I want the row element to be taken randomly from 1 to 9.
Where mask = False, I want the row element to be taken randomly from 1 to 3.

How can this be done efficiently? A sample output is shown below.

[[2, 7, 2, 9, 7, 1],
 [7, 2, 1, 7, 8, 3],
 [9, 1, 3, 5, 7, 3],
 [5, 7, 1, 4, 5, 2],
 [7, 7, 2, 6, 9, 1]]

Quang Hoang · Accepted Answer · 2020-11-18 23:42:05Z

1

One idea is to sample randomly between 0 to 1, then multiply with 9 or 3 depending on mask, and finally add 1 to move the sample.

rand = np.random.rand(5,len(mask))
is3 = (1-mask).astype(int)

# out is random from 0-8 or 0-2 depending on `is3`
out = (rand*val[is3]).astype(int)

# move out by `1`:
out = (out + 1)

Output:

array([[4, 9, 3, 6, 2, 1],
       [1, 8, 2, 7, 1, 3],
       [8, 2, 1, 2, 3, 2],
       [4, 3, 2, 2, 3, 2],
       [5, 8, 1, 5, 6, 1]])

edited Nov 18, 2020 at 23:42

answered Nov 18, 2020 at 22:54

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

user109387 Over a year ago

Very useful insight that I can use in other places!

Collectives™ on Stack Overflow

array rows where the random-integer elements may have different ranges

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related