是否可以使 numpy 阵列上的 function 更高效？

Question

这里a是 integer 索引的一维数组。 为了给出一些上下文， a是第一个分子的原子索引。 返回是n相同分子的原子索引，每个分子都包含step原子。 这个 function 基本上将相同的原子选择应用于许多分子

def f(a, step, n):
    a.shape=1,-1
    got = np.repeat(a, n, axis=0)
    blocks = np.arange(0, n*step, step)
    blocks.shape = n, -1
    return got + blocks

例如

In [52]: f(np.array([1,3,4]), 10, 4)
Out[52]:
array([[ 1,  3,  4],
       [11, 13, 14],
       [21, 23, 24],
       [31, 33, 34]])

Answer 1

看起来广播应该足够了：

def f(a, step, n):
    return a + np.arange(0, n*step, step)[:, None]

f(np.array([1,3,4]), 10, 4)

output：

array([[ 1,  3,  4],
       [11, 13, 14],
       [21, 23, 24],
       [31, 33, 34]])