Create random integer ndarray sampled from different span per element-CodePudding

I want to generate an ndarray a full of random integers which are sampled from different ranges according to another array span. For example:

import numpy as np

span = [5,6,7,8,9]

def get_a(span, count):
  a = np.stack([np.random.choice(i, count) for i in span], axis=0)
  return a

get_a(span,2)

Is there a fast way to do get_a?

CodePudding user response：

Yes. Yours:

import timeit
import numpy as np

span = np.arange(1,100)

def get_a(span, count):
  a = np.stack([np.random.choice(i, count) for i in span], axis=0)
  return a

%timeit get_a(span,2)

2.32 ms ± 254 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

My solution is 100s times faster for largish arrays:

def get_b(span, count):
    b = (np.random.rand(len(span), count)*span[:,None]).astype(int)
    return b

%timeit get_b(span,2)

6.91 µs ± 267 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)