优化numpy二维数组列高度统计的for循环 适配高频调用性能需求

核心优化思路 #

完全替换原生Python嵌套循环，所有逻辑通过NumPy内置的C级向量化操作完成，单帧计算过程无Python层循环跳转，适配高频调用的性能要求。

分步实现逻辑 #

• 跳过棋盘前2行的统计区域切片
  直接对原数组做无拷贝切片，定位有效统计范围：
  valid_area = self.board[2:, :]
• 逐列定位第一个非零块的位置
  生成有效区域的非零掩码，沿列方向（从上到下）查找第一个非零值的位置，不需要逐列遍历：

  # 生成非零值掩码，值为非0的位置标记为True
  nonzero_mask = valid_area != 0
  # 逐列查找第一个True的位置（即第一个非零值在有效区域内的行索引）
  first_pos_in_valid = nonzero_mask.argmax(axis=0)
  # 标记每列是否存在非零值，处理全空列的边界情况
  has_block = nonzero_mask.any(axis=0)
  

• 按规则计算每列的实际高度
  结合切片的行偏移量（前2行）还原原棋盘的非零行索引，套用高度计算公式：

  # 还原原棋盘上第一个非零块的行索引
  orig_row_idx = first_pos_in_valid + 2
  # 按规则计算列高度，全空列高度设为0
  col_heights = np.where(has_block, self.rows - orig_row_idx, 0)
  

• 累加相邻列高度的绝对差值
  用NumPy内置的差分计算接口直接得到相邻列的高度差，取绝对值后求和：
  total_result = np.abs(np.diff(col_heights)).sum()

效果校验 #

针对示例给出的列高度序列[15,15,19,20,14,19,20,20,20,20]，上述计算逻辑输出的累加差值为|15-15| + |15-19| + |19-20| + |20-14| + |14-19| + |19-20| + |20-20| + |20-20| + |20-20| = 17，和手动计算结果完全一致。

注意避坑点 #

• 轴参数设置：逐列查找非零值时，argmax的axis必须设为0，若误设为1会变成逐行遍历，完全偏离统计逻辑。
• 全空列边界处理：若某列跳过前2行后全为零值，argmax会默认返回索引0，必须通过has_block的判断将这类列的高度置为0，否则会出现高度计算虚高的错误。
• 内存优化：所有切片操作默认返回原数组的视图，不会触发全量数组拷贝，高频调用时内存开销极低。

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