NumPy数组逐元素自定义运算：循环实现及大规模数组高效处理方案咨询

对应实现代码如下：

import numpy as np
arr = np.array([[1, 2], [4, 5]], np.int32)
squarred = np.zeros_like(arr)
# 逐行遍历处理
for row_idx in range(arr.shape[0]):
    squarred[row_idx, 0] = arr[row_idx, 0] ** 2
    squarred[row_idx, 1] = arr[row_idx, 1] ** 3

这种写法逻辑直观，但仅适合小规模数组使用：Python原生循环有较高的执行开销，数组规模越大，运行速度会比你原有的矢量化写法慢越多，万行以上的数组性能差距可达数十倍。

按实际场景的优先级选择即可：

• 首选NumPy原生矢量化实现：如果你的自定义变换逻辑可以拆解为NumPy内置的算术运算、广播、数学函数的组合，直接用矢量化写法是性能最高的，完全规避Python层循环开销。
  以你示例的需求为例，还可以简化为更紧凑的写法，适配任意行数的2列数组：

  squarred = np.column_stack([arr[:,0] ** 2, arr[:,1] ** 3])
  

• 复杂自定义逻辑优先用Numba JIT编译：如果变换逻辑非常特殊，没法用NumPy内置操作拼接实现，推荐用Numba的JIT装饰器编译自定义循环函数，执行效率接近原生C语言，完美适配大规模数组。
  通用的逐元素自定义变换实现模板如下：

  import numpy as np
  from numba import jit
  
  @jit(nopython=True) # 开启无Python对象模式，最大化性能
  def custom_arr_transform(input_arr):
      # 输出和输入维度完全一致
      output_arr = np.zeros_like(input_arr)
      row_cnt, col_cnt = input_arr.shape
      for i in range(row_cnt):
          for j in range(col_cnt):
              # 此处可替换为任意自定义变换逻辑
              if j == 0:
                  output_arr[i,j] = input_arr[i,j] ** 2
              elif j == 1:
                  output_arr[i,j] = input_arr[i,j] ** 3
              # 可扩展更多列的自定义规则
      return output_arr
  
  # 调用方法
  squarred = custom_arr_transform(arr)
  

• 补充说明：np.vectorize、np.apply_along_axis这类方法本质是Python层循环的语法糖，性能远不如上述两种方案，仅适合临时处理小数组时简化代码使用，不推荐大规模场景使用。

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