已知传入Python函数的数组固定为3x1大小，如何指定大小优化计算效率？

首先你之前写的def FUNC(p0=[3])的语法完全达不到指定入参长度的效果，这个写法是给p0设置默认值为[3]这个长度为1的列表，而且Python可变默认值存在隐藏坑，非常不推荐这么写。

关于指定数组大小能不能节省耗时，分两种场景说明：

• 如果你用的是原生Python列表：原生列表没有编译期长度优化机制，就算你额外加运行时长度校验，也只会多出来检查的开销，不会减少计算耗时，除非你需要做入参合法性校验避免后续逻辑报错，否则没必要做。
• 如果你用的是NumPy数组（3x1数组的数值计算场景一般都会用NumPy）：固定形状的数组确实可以带来性能提升，NumPy可以利用固定形状信息做内存预分配、SIMD指令优化，还能省去运行时动态判断数组形状的开销。这种场景可以用类型标注配合numpy.typing指定数组形状，需要运行时校验的话可以加断言，示例代码如下：

import numpy as np
import numpy.typing as npt

# 标注p0是float64类型的数组，调用时固定传(3,1)形状即可触发优化
def FUNC(p0: npt.NDArray[np.float64]) -> None:
    # 仅需要做入参合法性校验时加这行，优化本身不需要
    assert p0.shape == (3, 1), "入参必须为3x1数组"
    print(p0[1, 0])

# 调用示例
FUNC(np.array([[1], [1], [2]]))

• 如果你追求极致性能，可以搭配Numba的JIT编译使用，指定数组形状后Numba会生成高度优化的机器码，性能会比纯Python或者原生NumPy有数量级的提升，示例如下：

from numba import jit
import numpy as np
import numpy.typing as npt

@jit(nopython=True)
def FUNC(p0: npt.NDArray[np.float64]) -> float:
    return p0[0, 0] + p0[1, 0] + p0[2, 0]

如果你的场景只用原生Python列表，那这部分性能差异完全可以忽略，不需要做额外优化。

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