如何用Python优雅计算三角网格的顶点距离矩阵？

问题背景 #

我用Numpy数组表示三角网格，有两个核心矩阵：

• coordinates：形状为3×n或n×3的矩阵（示例中为n×3，每行对应一个顶点的三维坐标；若为3×n则每列对应一个顶点）
• connectivity：形状为n×n的0-1矩阵，connectivity[i,j] == 1表示顶点i和j连通

需要生成一个n×n的distances矩阵，仅对连通位置计算顶点间的欧氏距离，非连通位置设为0（或保持初始值）。

优雅实现方案 #

利用Numpy的广播机制和矩阵运算，避免显式循环，高效完成计算：

import numpy as np

# 示例输入
coordinates = np.array(
    [
        [-1, -1, 0],    # A
        [1, -1, 0],     # B
        [1, 1, 0],      # C
        [-1, 1, 0]      # D
    ], 
    dtype=np.float32
)

connectivity = np.array(
    [
        [0, 1, 1, 1],    # A-B, A-C, A-D
        [1, 0, 1, 0],    # B-A, B-C
        [1, 1, 0, 1],    # C-A, C-B, C-D
        [1, 0, 1, 0],    # D-A, D-C
    ],
    dtype=np.int32
)

# 若coordinates是3×n的格式，先转置为n×3
# coordinates = coordinates.T

# 步骤1：计算所有顶点对的欧氏距离矩阵
# 广播计算坐标差：(n,1,3) - (1,n,3) → (n,n,3)
coord_diff = coordinates[:, None, :] - coordinates[None, :, :]
# 对每个顶点对的坐标差求欧氏距离，得到(n,n)的距离矩阵
all_distances = np.linalg.norm(coord_diff, axis=2)

# 步骤2：用连通性矩阵过滤，仅保留连通位置的距离，其余设为0
distances = np.where(connectivity == 1, all_distances, 0.0)

# 验证结果（保留3位小数对齐示例）
print(np.round(distances, 3))

代码说明 #

1. 坐标格式适配：如果你的coordinates是3×n的列优先存储格式，只需先转置为n×3，方便后续广播计算。
2. 广播计算坐标差：通过维度扩展，将原坐标数组转为(n,1,3)和(1,n,3)的形状，相减后得到(n,n,3)的坐标差数组，每个元素对应一对顶点的坐标差值。
3. 批量计算欧氏距离：np.linalg.norm直接对所有顶点对的三维坐标差求范数，一步生成全量距离矩阵，无需循环遍历。
4. 掩码过滤结果：np.where根据连通性矩阵的取值，自动保留连通顶点的距离，非连通位置设为0，完全匹配需求。

输出结果 #

运行上述代码后，输出与示例预期一致：

[[0.    2.    2.828 2.   ]
 [2.    0.    2.    0.   ]
 [2.828 2.    0.    2.   ]
 [2.    0.    2.    0.   ]]

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