以高效的方式使用不同参数调用函数来修改numpy数组

Question

我想从这段代码中消除低效的for循环。

import numpy as np

x = np.zeros((5,5))

for i in range(5):
    x[i] = np.random.choice(i+1, 5)

在保持给定输出的同时

[[0. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0.]
 [0. 2. 2. 1. 0.]
 [1. 2. 3. 1. 0.]
 [1. 0. 3. 3. 1.]]

我已经尝试过这个。

i = np.arange(5)
x[i] = np.random.choice(i+1, 5)

但它输出的是：

[[0. 1. 1. 3. 3.]
 [0. 1. 1. 3. 3.]
 [0. 1. 1. 3. 3.]
 [0. 1. 1. 3. 3.]
 [0. 1. 1. 3. 3.]]

有没有可能去掉循环？如果不行，对于一个大数组和许多重复的情况，最有效的处理方式是什么？

- Xavi Reyes

Answer 1

创建一个随机整数数组，每行的最高数字作为列数。因此，我们可以使用 np.random.randint 函数，并将其 high 参数设置为列数。然后，执行模运算以在每行上设置不同的限制，这些限制由行号定义。因此，我们将得到以下向量化实现 -

def create_rand_limited_per_row(m,n):
    s = np.arange(1,m+1)
    return np.random.randint(low=0,high=n,size=(m,n))%s[:,None]

样例运行 -

In [45]: create_rand_limited_per_row(m=5,n=5)
Out[45]: 
array([[0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [1, 2, 0, 2, 1],
       [0, 0, 1, 3, 0],
       [1, 2, 3, 3, 2]])

为了利用numexpr模块的多核支持处理大数据 -

import numexpr as ne

def create_rand_limited_per_row_numepxr(m,n):
    s = np.arange(1,m+1)[:,None]
    a = np.random.randint(0,n,(m,n))
    return ne.evaluate('a%s')

Benchmarking

# Original approach
def create_rand_limited_per_row_loopy(m,n):
    x = np.empty((m,n),dtype=int)
    for i in range(m):
        x[i] = np.random.choice(i+1, n)
    return x

1k x 1k数据的时序 -

In [71]: %timeit create_rand_limited_per_row_loopy(m=1000,n=1000)
10 loops, best of 3: 20.6 ms per loop

In [72]: %timeit create_rand_limited_per_row(m=1000,n=1000)
100 loops, best of 3: 14.3 ms per loop

In [73]: %timeit create_rand_limited_per_row_numepxr(m=1000,n=1000)
100 loops, best of 3: 6.98 ms per loop

- Divakar