能否优化这个基于numpy数组的原子索引生成函数？

• 参数说明：
  • a：一维整数数组，代表第一个分子的原子索引
  • step：单个分子包含的原子数量
  • n：需要生成索引的分子总数

该函数用于将针对单个分子的原子选择逻辑，批量应用到n个相同分子上，生成所有分子对应的目标原子索引。

实现代码 #

import numpy as np

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]])

内容的提问来源于stack exchange，提问作者nos