Numpy快速合并同类型列：含重复目标点的距离矩阵构建

给定Numpy数组A，第一列记录目标点坐标，第二列记录对应距离。需构建距离矩阵D：

• 行对应源点（即数组A的每一行）
• 列对应唯一目标点坐标
• 未知距离用np.inf填充

直接使用np.unique获取唯一列时，会丢失重复目标点对应的其他距离信息（如示例中第5行的0.91、第8行的0.82均丢失），现需纯Numpy高效实现（数据量超1万对点，禁止使用Pandas）。

示例代码及问题 #

import numpy as np

A = np.array([[7.3, 0.25],
              [2.5, 0.32],
              [3.7, 0.45],
              [6.1, 0.55],
              [2.5, 0.91],
              [4.8, 0.77],
              [8.6, 0.35],
              [6.1, 0.82]])
D = np.ones((A.shape[0],A.shape[0])) * np.inf
np.fill_diagonal(D, A[:,1])
unq_ids = np.sort(np.unique(A[:,0], return_index=True)[1])
D = D[:,unq_ids]
print(D)

错误输出 #

array([[0.25,  inf,  inf,  inf,  inf,  inf],
       [ inf, 0.32,  inf,  inf,  inf,  inf],
       [ inf,  inf, 0.45,  inf,  inf,  inf],
       [ inf,  inf,  inf, 0.55,  inf,  inf],
       [ inf,  inf,  inf,  inf,  inf,  inf],
       [ inf,  inf,  inf,  inf, 0.77,  inf],
       [ inf,  inf,  inf,  inf,  inf, 0.35],
       [ inf,  inf,  inf,  inf,  inf,  inf]])

期望输出 #

array([[0.25,  inf,  inf,  inf,  inf,  inf],
       [ inf, 0.32,  inf,  inf,  inf,  inf],
       [ inf,  inf, 0.45,  inf,  inf,  inf],
       [ inf,  inf,  inf, 0.55,  inf,  inf],
       [ inf, 0.91,  inf,  inf,  inf,  inf],
       [ inf,  inf,  inf,  inf, 0.77,  inf],
       [ inf,  inf,  inf,  inf,  inf, 0.35],
       [ inf,  inf,  inf, 0.82,  inf,  inf]])

解决方案 #

利用np.unique的return_inverse参数获取每个原始目标点对应的唯一列索引，再通过np.put_along_axis批量填充距离，全程无循环，适合大数据量：

import numpy as np

A = np.array([[7.3, 0.25],
              [2.5, 0.32],
              [3.7, 0.45],
              [6.1, 0.55],
              [2.5, 0.91],
              [4.8, 0.77],
              [8.6, 0.35],
              [6.1, 0.82]])

# 1. 获取唯一目标点，以及每个原始目标点对应的唯一列索引
unq_targets, col_indices = np.unique(A[:, 0], return_inverse=True)

# 2. 初始化距离矩阵：行数为源点总数，列数为唯一目标点数量
D = np.full((A.shape[0], len(unq_targets)), np.inf)

# 3. 批量填充距离：将每个源点的距离放到对应唯一目标点的列位置
np.put_along_axis(D, col_indices.reshape(-1, 1), A[:, 1].reshape(-1, 1), axis=1)

print(D)

方案说明 #

• return_inverse=True返回的col_indices数组，每个元素对应原始目标点在唯一目标点数组中的索引，完美映射重复目标点到同一列
• np.put_along_axis是Numpy内置的批量赋值函数，效率远高于循环，处理1万级数据毫无压力
• 最终矩阵的列顺序与unq_targets的排序一致，符合常规需求

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