Python函数性能优化咨询：多次调用average_over_neighbours效率低下

我编写了如下函数：

def average_over_neighbours(
    distances: npt.NDArray[np.float32],
    velocities: npt.NDArray[np.float32],
    n_neighbours: int = 3,
    weights: npt.NDArray[np.float32] = np.array([]),
) -> tuple[int, npt.NDArray[np.float32]]:

    in_range = distances < search_range
    neighbours: npt.NDArray[np.int32] = np.nonzero(in_range)[0][
        distances[in_range].argsort(kind="mergesort")
    ][:n_neighbours]

    if neighbours.size == 0:
        return 0, np.array([0.0, 0.0])

    velocities = velocities[neighbours]
    weights = weights[neighbours] if weights.size != 0 else None
    average_velocity: npt.NDArray[np.float32] = np.average(
        velocities, weights=weights, axis=0
    )
    return len(neighbours), average_velocity

需要在for循环中多次调用该函数：

for point in my_points:
    average_over_neighbours(distances, velocities)

当前程序性能极差，想明确以下问题：

• 性能瓶颈是函数本身存在低效点，还是仅因数千次调用的开销导致？
• 将循环移至函数内部，能否借助减少Python函数调用开销显著提升性能？
• 曾尝试重构但未看到明显改善，又担心重构降低代码可读性，寻求可行的优化方案。

一、函数本身的低效点排查与优化 #

1. 排序方式冗余
   你使用了mergesort稳定排序，但仅需取前n_neighbours个最小距离的邻居，全量稳定排序完全没必要。换成quicksort（argsort默认方式）或改用np.argpartition能大幅降低时间复杂度：

   • 全量排序是O(m log m)，argpartition是O(m)，对于大尺寸distances数组，前者耗时会显著更高。
   • 优化后的邻居查找逻辑示例：

     in_range_idx = np.where(distances < search_range)[0]
     if not in_range_idx.size:
         return 0, np.array([0.0, 0.0])
     # 直接找到前n_neighbours个最小元素的索引，无需全量排序
     partition_idx = np.argpartition(distances[in_range_idx], n_neighbours-1)[:n_neighbours]
     neighbours = in_range_idx[partition_idx]
     

2. 默认参数与判断冗余
   函数默认weights为空数组，每次调用都会创建空数组实例，且每次都要判断weights.size !=0。可以将默认参数改为weights=None，减少不必要的数组创建与判断开销：

   def average_over_neighbours(
       distances: npt.NDArray[np.float32],
       velocities: npt.NDArray[np.float32],
       n_neighbours: int = 3,
       weights: npt.NDArray[np.float32] = None,
   ) -> tuple[int, npt.NDArray[np.float32]]:
       # 内部判断改为
       weights = weights[neighbours] if weights is not None else None
   

二、函数调用开销的影响 #

Python单次函数调用开销约几十纳秒，数千次调用的总开销通常不会成为核心瓶颈——你的函数内部包含排序、数组切片等耗时操作，这些才是性能差的主要原因。

单纯将循环移到函数内部，若仍用Python循环逐个处理，性能提升非常有限；但如果改成批量向量化处理（比如将distances设为二维数组，每一行对应一个point的距离数据），用numpy的向量化操作替代Python循环，才能真正大幅提升性能。

三、兼顾性能与可读性的优化方案 #

1. 批量向量化处理
   如果my_points对应的distances可以组织成二维数组（shape=(N, M)，N为点数，M为候选邻居数），可以用numpy的广播、按行argpartition等操作，一次性完成所有点的邻居平均速度计算，完全避免Python循环。

2. 用Numba编译加速
   如果无法调整数据结构实现批量处理，可用numba.jit将函数编译为机器码，既保留原代码的可读性，又能消除Python循环与函数调用的开销。示例：

   from numba import jit, float32, int32
   import numpy as np
   import numpy.typing as npt
   
   @jit(tuple((float32[:], float32[:, :], int32, float32[:])), nopython=True)
   def average_over_neighbours(
       distances: npt.NDArray[np.float32],
       velocities: npt.NDArray[np.float32],
       n_neighbours: int = 3,
       weights: npt.NDArray[np.float32] = None,
   ) -> tuple[int, npt.NDArray[np.float32]]:
       in_range = distances < search_range
       in_range_idx = np.where(in_range)[0]
       if in_range_idx.size == 0:
           return 0, np.array([0.0, 0.0], dtype=np.float32)
       
       # 用argpartition获取前n个最小元素索引
       partition_idx = np.argpartition(distances[in_range_idx], n_neighbours-1)[:n_neighbours]
       neighbours = in_range_idx[partition_idx]
       
       velocities_neigh = velocities[neighbours]
       if weights is not None and weights.size > 0:
           weights_neigh = weights[neighbours]
           sum_weights = np.sum(weights_neigh)
           if sum_weights == 0:
               return len(neighbours), np.array([0.0, 0.0], dtype=np.float32)
           avg_vel = np.sum(velocities_neigh * weights_neigh[:, None], axis=0) / sum_weights
       else:
           avg_vel = np.mean(velocities_neigh, axis=0)
       return len(neighbours), avg_vel
   

总结 #

• 函数内部的全量排序是核心低效点，替换为np.argpartition可大幅降低时间复杂度；
• 数千次函数调用不是性能瓶颈，批量向量化处理或Numba编译才是有效的性能提升手段；
• 优先选择Numba编译（代码改动小、可读性保留好），若数据结构允许则用批量向量化操作。

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