Python确定性k-means算法不同初始化函数性能测试及优化咨询

我很理解你现在的困扰——用一堆if-elif来管理不同的初始化方法确实不够灵活，而且后续加新方法会越来越乱。下面给你几个实用的优化方向，以及性能测试的具体方法：

最常用也最优雅的方式是用字典映射来替代条件分支，把初始化方法名称和对应的函数直接绑定，这样不仅代码更整洁，后续新增初始化方法时只需要往字典里加条目就行，不用修改条件判断逻辑。

举个具体的实现例子：

# 先定义一个初始化方法的映射字典
init_methods = {
    'RANDOM': lambda unique_datap, k: unique_datap[np.random.choice(unique_datap.shape[0], k, replace=False)].astype(np.float32),
    'FFT': lambda unique_datap, el_count, k, rgb_distance: deterministic_fft(unique_datap, el_count, k, rgb_distance).astype(np.float32),
    'UMDI': lambda unique_datap, el_count, k, rgb_distance: uniform_mode_dist_init(unique_datap, el_count, k, rgb_distance).astype(np.float32)
}

# 然后调用的时候直接从字典取
try:
    # 注意：如果RANDOM方法不需要el_count和rgb_distance，可以调整lambda参数或者统一函数接口
    c_means = init_methods[init](unique_datap, el_count, k, rgb_distance)
except KeyError:
    raise ValueError(f'Unknown initialization method: {init}')

如果你的各个初始化函数参数不完全一致，也可以稍微调整字典的结构：比如给每个函数适配统一的参数接口，或者把参数打包成kwargs传入，这样字典的调用逻辑会更统一。

接下来聊聊你关心的执行时间测量、内存占用监控，以及结果保存：

1. 测量执行时间 #

• 简单的单次计时可以用time模块：

  import time
  
  start_time = time.time()
  # 运行你的k-means初始化+完整聚类逻辑
  c_means = init_methods[init](unique_datap, el_count, k, rgb_distance)
  # 后续聚类迭代步骤...
  end_time = time.time()
  print(f"Total execution time: {end_time - start_time:.4f} seconds")
  

• 如果需要更准确的多次平均计时，推荐用timeit模块，它会自动排除部分系统干扰：

  import timeit
  
  def run_kmeans_init():
      return init_methods[init](unique_datap, el_count, k, rgb_distance)
  
  # 运行100次取平均结果
  avg_time = timeit.timeit(run_kmeans_init, number=100) / 100
  print(f"Average initialization time: {avg_time:.6f} seconds")
  

2. 监控内存占用 #

• 对于基础的内存使用情况，可以用Python标准库tracemalloc：

  import tracemalloc
  
  tracemalloc.start()
  # 运行初始化逻辑
  c_means = init_methods[init](unique_datap, el_count, k, rgb_distance)
  current, peak = tracemalloc.get_traced_memory()
  print(f"Current memory usage: {current / 10**6:.2f} MB")
  print(f"Peak memory usage: {peak / 10**6:.2f} MB")
  tracemalloc.stop()
  

• 如果需要更细粒度的内存分析（比如函数级别的内存消耗），可以用第三方库memory_profiler：安装后用装饰器@profile标记要分析的函数，运行时就能看到逐行的内存使用报告。

3. 保存算法输出 #

• 保存聚类中心（c_means）可以用numpy的内置方法，比如保存为NPY文件：

  import numpy as np
  np.save(f'kmeans_centers_{init}_k{k}.npy', c_means)
  

• 如果需要保存更完整的实验结果（比如每个样本的聚类标签、执行时间、内存数据等），可以用pickle序列化整个结果对象：

  import pickle
  
  result = {
      'init_method': init,
      'k_value': k,
      'cluster_centers': c_means,
      'total_execution_time': end_time - start_time,
      'peak_memory_mb': peak / 10**6
  }
  
  with open(f'kmeans_result_{init}_k{k}.pkl', 'wb') as f:
      pickle.dump(result, f)
  

• 要是需要可读性更强的格式（比如给非技术人员查看），可以把关键结果导出为CSV：

  import pandas as pd
  # 把聚类中心转成DataFrame保存
  pd.DataFrame(c_means).to_csv(f'kmeans_centers_{init}_k{k}.csv', index=False)
  

最后，这种字典映射的方式在Python工程化代码里非常常见，既能提升代码的可维护性，也符合"开放封闭原则"——新增功能不用修改原有逻辑，只需要扩展新的条目。

内容的提问来源于stack exchange，提问作者Juan Eduardo Casilla Camarillo