如何高效查找两点连线上3D numpy数组的对应销售数据点

前置说明 #

你的数据规模为80层 * 125*125平面栅格，且所有层的起点A、终点B的xy坐标完全一致，完全不需要逐层重复计算连线路径，核心优化思路是单次计算2维栅格连线，批量映射到所有层，可以把计算量降到逐层计算的1/80。

具体实现步骤 #

• 统一索引映射规则
  3维销售额数组的维度顺序固定为(层号z, y轴坐标, x轴坐标)时，单元格编号和3维索引可以直接通过公式换算，不需要额外建映射表，计算速度极快：
  • 已知索引(z,y,x)，对应单元格编号：cell_id = z * 125 * 125 + y * 125 + x
  • 已知单元格编号，对应3维索引：z = cell_id // (125*125); y = (cell_id % (125*125)) // 125; x = cell_id % 125
• 单次计算2维平面的连线栅格点
  直接调用栅格画线的成熟底层实现，不要自己写纯Python循环的画线逻辑：
  推荐用skimage.draw.line，是Bresenham画线算法的C级实现，输入A、B点的坐标，一次返回连线上所有栅格点的y、x坐标数组，125*125尺寸下计算耗时小于1ms。

  注意传参顺序：该函数输入顺序为(y0, x0, y1, x1)，不要把x、y顺序传反导致坐标错位。

• 广播生成所有层的全量索引
  拿到单层连线的坐标数组后，用numpy广播机制一次性生成所有80层的索引，全程无Python层循环：
  • 层号数组构造为形状(80, 1)的列向量
  • 单层连线的y、x数组形状为(n,)（n为单条连线的点总数），直接广播为(80, n)的矩阵
  • 用三个索引矩阵直接取3维数组的值，同时批量换算所有点的单元格编号即可。

可直接运行的代码示例 #

import numpy as np
from skimage.draw import line

# 固定参数
GRID_SIZE = 125
LAYER_NUM = 80
# 替换为你实际的A、B点x、y坐标
A_X, A_Y = 22, 47
B_X, B_Y = 103, 96

# 从DataFrame转换得到3维销售额数组，形状为(80,125,125)
# sales_3d = df["sales"].to_numpy().reshape(LAYER_NUM, GRID_SIZE, GRID_SIZE)

# 1. 单次计算2维平面连线点
y_coords, x_coords = line(A_Y, A_X, B_Y, B_X)
point_per_layer = len(y_coords)

# 2. 广播生成全量3维索引
z_index = np.arange(LAYER_NUM)[:, None]
y_index = np.broadcast_to(y_coords, (LAYER_NUM, point_per_layer))
x_index = np.broadcast_to(x_coords, (LAYER_NUM, point_per_layer))

# 3. 提取所有连线上的销售额，形状为(80, point_per_layer)
line_sales_data = sales_3d[z_index, y_index, x_index]

# 4. 批量计算所有点对应的单元格编号，展平为1维数组
all_line_cell_ids = (z_index * GRID_SIZE * GRID_SIZE + y_index * GRID_SIZE + x_index).ravel()

性能说明 #

• 全流程核心计算均为C级实现的numpy、scikit-image算子，无Python循环，总耗时在普通PC上小于10ms，比逐行遍历DataFrame、逐层循环计算的方案快100倍以上。
• 如果不想引入scikit-image依赖，也可以自行实现Bresenham画线算法，纯Python实现的单平面连线计算耗时约10~20ms，对于你的数据规模也完全够用。

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