用Numpy替换弹跳球模拟中交换速度向量的For循环以提升FPS

你的原嵌套循环属于O(n²)复杂度，不仅会重复检测同一对球（比如i=0,j=1和i=1,j=0会做两次判断），还因逐个操作Python对象带来额外开销，是FPS上不去的核心原因。用Numpy的向量化操作能把这些循环转为底层C实现的批量计算，大幅提升性能，具体优化思路和代码如下：

1. 统一存储小球数据为Numpy数组 #

放弃单个ball对象的零散存储，把所有球的位置、速度统一转成Numpy数组，方便批量操作：

import numpy as np

# 初始化时把现有ball数据转成数组
ball_count = len(balls)
positions = np.array([ball.xy for ball in balls])  # shape: (球数量, 2)，每行对应(x,y)
velocities = np.array([ball.v for ball in balls])   # shape: (球数量, 2)，每行对应(vx, vy)

2. 批量检测碰撞并交换速度 #

用Numpy广播机制计算所有球对的距离，筛选碰撞对后批量交换速度：

def update_with_numpy():
    global positions, velocities
    # 计算所有球对的位置差，利用广播生成(n,n,2)的矩阵
    pos_diff = positions[:, None] - positions
    # 计算距离平方（避免开根号，减少计算开销）
    dist_sq = np.sum(pos_diff ** 2, axis=2)
    # 生成碰撞掩码：距离平方小于15²=225，且排除球自身和重复配对
    collision_mask = (dist_sq < 225) & (np.eye(ball_count, dtype=bool) == False)
    # 只取上三角的碰撞对，每个碰撞仅处理一次
    i, j = np.where(np.triu(collision_mask, k=1))
    
    # 批量交换碰撞球对的速度
    velocities[[i, j]] = velocities[[j, i]]
    
    # 若需要保留原ball对象，同步更新速度（不需要则可省略此循环）
    for idx in range(ball_count):
        balls[idx].v = velocities[idx].tolist()

3. 额外优化细节 #

• 避免重复计算：用np.triu只处理i<j的球对，直接砍掉一半的碰撞检测量
• 跳过开根号操作：比较距离平方和225的大小，结果和直接比较距离完全一致，却能省去大量耗时的开根号计算
• 减少Python对象交互：核心计算全在Numpy数组上完成，仅在必要时同步回原对象，把Python层面的循环开销降到最低

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