Python Numpy如何高效按不规则索引将小矩阵行复制到大矩阵对应行

你之前使用的A[3,3,2,2,3,1]写法属于numpy高级索引，对于非连续的索引选取，底层走通用拷贝逻辑，且每次执行都会重新申请结果数组的内存空间，因此性能较差，甚至不如手动实现的逐行for循环。针对映射规则固定的场景，可以通过以下方案大幅提升性能：

实现步骤 #

• 提前定义固定的行映射索引（注意numpy默认使用0-based索引，需要将你示例的1-based索引减1），同时预分配矩阵B的内存，避免每次运算重复申请空间：

  import numpy as np
  
  # 示例映射规则[3,3,2,2,3,1]转0-based，实际使用时替换为你对应的2000个元素的映射数组即可
  map_idx = np.array([2,2,1,1,2,0], dtype=np.intp)
  # 预分配和目标尺寸一致的B矩阵，dtype和A保持一致
  B = np.empty((2000, 500), dtype=A.dtype)
  

• 每次矩阵A动态更新后，使用np.take接口原地更新B，该接口针对沿轴选取元素的场景做了底层优化，直接写入预分配的B内存空间，省去了额外的内存拷贝开销：

  np.take(A, map_idx, axis=0, out=B)
  

额外优化点 #

• 如果你的矩阵A不是行优先（C连续）存储，提前转为连续数组可以进一步提升取行速度：

  A = np.ascontiguousarray(A)
  

• 若映射规则存在大量连续重复的行，可以提前统计每个行的重复次数，用np.repeat实现，性能和np.take接近，代码可读性更高：

  # 示例：行2重复3次、行1重复2次、行0重复1次，按实际规则调整即可
  np.repeat(A[[2,1,0]], [3,2,1], axis=0, out=B)
  

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