如何为Python数组行分配索引 快速调用预计算BD距离矩阵

你遇到的两个核心问题：

1. 固定数组的行向量无法和索引绑定，导致预计算好的距离矩阵无法快速查表
2. 聚类迭代过程中重复计算BD距离，运行耗时过高

1. 给固定数组行绑定索引的实现方法 #

针对内容不变的原始数组（比如你的psnr_bitrate样本集），两种方法可以快速关联行和索引：

• 遍历数组时直接用enumerate()同步获取索引和行值，不要单独把行抽出来之后再反查索引：

  for idx, row in enumerate(A):
      # 这里idx就是当前行的索引，row是行内容，可以直接用来查询距离矩阵
  

• 如果必须根据拿到的行值反查索引，提前构建行到索引的映射字典，查询复杂度为O(1)：

  # 数组构建完成后执行一次即可
  row_to_idx = {tuple(row): idx for idx, row in enumerate(A)}
  
  # 拿到任意行A1、A2时，直接查询索引取距离
  i = row_to_idx[tuple(A1)]
  j = row_to_idx[tuple(A2)]
  d_ij = D[i,j]
  

  注意：numpy数组是可变对象，不能作为字典键，所以转成不可变的元组再存；这个方法仅适用于内容不会修改的固定数组。

2. 聚类过程的距离计算优化 #

首先明确你用的聚类类型，两种场景优化方式不同：

场景1：质心始终来自原始样本（K-medoids聚类） #

这种场景下所有要计算距离的点都在原始600个样本里，完全可以一次性预计算全局距离矩阵，后续迭代全程查表，不需要重复调用BD计算函数，速度可以提升几十到上百倍：

1. 提前一次性计算600x600的对称距离矩阵，利用距离对称性减少一半计算量：

   m = psnr_bitrate.shape[0]
   wr = 0.5
   wq = 0.5
   D = np.zeros((m, m))
   
   for i in range(m):
       for j in range(i, m):
           BD_R = bd_rate(rate, psnr_bitrate[i], rate, psnr_bitrate[j])
           BD_R = np.inf if BD_R == -2 else BD_R
           BD_Q = bd_PSNR(rate, psnr_bitrate[i], rate, psnr_bitrate[j])
           BD_Q = np.inf if BD_Q == -2 else BD_Q
           dist = np.abs(wr * BD_R + wq * BD_Q)
           D[i,j] = dist
           D[j,i] = dist
   

2. 迭代过程中只需要保存每个质心对应的原始数组索引，不需要存质心的数值本身，所有距离直接从D中切片获取，完全跳过重复的BD计算：

   def kmedoids_BD(psnr_bitrate, K, initial_centroid_ids, n_itr=10):
       m = psnr_bitrate.shape[0]
       centroid_ids = initial_centroid_ids.copy()
       for itr in range(n_itr):
           # 直接切片取所有样本到当前K个质心的距离，无额外计算
           BD = D[:, centroid_ids]
           labels = np.argmin(BD, axis=1)
   
           # 更新质心：选每个簇内平均距离最小的样本点
           new_centroid_ids = np.zeros(K, dtype=int)
           for k in range(K):
               cluster_members = np.where(labels == k)[0]
               if len(cluster_members) == 0:
                   new_centroid_ids[k] = np.random.randint(0, m)
                   continue
               cluster_dist_sum = D[cluster_members][:, cluster_members].sum(axis=1)
               new_centroid_ids[k] = cluster_members[np.argmin(cluster_dist_sum)]
           
           # 收敛则提前退出
           if np.all(centroid_ids == new_centroid_ids):
               break
           centroid_ids = new_centroid_ids
       return centroid_ids, labels
   

场景2：质心是簇内样本均值（标准Kmeans聚类） #

这种场景下每次更新的质心是新生成的向量，不在原始样本集中，无法用预计算的全局距离矩阵查表，可以做两个优化减少耗时：

• 每次迭代只计算新生成的K个质心到所有样本的距离，不要重复计算样本之间的距离
• 尽量用numpy向量化操作替换多层Python原生循环，减少解释器开销

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