如何在Python中高效计算2097152个星系点间的距离？

原代码的核心问题 #

• 时间复杂度为O(n²)：当n=2097152时，需要计算约2×10¹²个点对，这在现有硬件下完全无法在合理时间内完成。
• np.append的低效性：每次调用都会重新分配数组内存并复制数据，进一步加剧性能损耗。
• （小问题）距离计算中的笔误：(z1-z2)**2不影响结果，但规范写法应为(z2-z1)**2（平方后结果一致）。

可行的高效解决方案 #

由于200多万个点的全量两两距离计算在存储和时间上都不现实，结合你提到的「特定尺度球壳」场景，我们可以通过空间索引+邻域查询的方式，只计算符合距离范围的点对，再统计分布。

步骤1：准备数据并构建KDTree #

KDTree是一种高效的空间索引结构，能快速查询指定半径内的邻居点：

import numpy as np
from scipy.spatial import KDTree

# 将x/y/z组合成坐标矩阵（shape: [2097152, 3]）
coords = np.vstack((x, y, z)).T

# 构建KDTree
tree = KDTree(coords)

步骤2：查询指定距离范围内的邻居 #

假设你关注的球壳尺度对应的最大距离为max_r，批量查询每个点在该范围内的邻居及距离：

max_r = 你的球壳最大半径  # 替换为实际数值，比如100

# query_ball_point返回每个点的邻居索引和对应距离
neighbor_distances, neighbor_indices = tree.query_ball_point(
    coords, r=max_r, return_distance=True
)

步骤3：收集非重复的点对距离 #

为避免重复计算（i<j和j<i的距离是同一个），只保留索引大于当前点的邻居距离：

all_valid_distances = []
for idx, (dists, idxs) in enumerate(zip(neighbor_distances, neighbor_indices)):
    # 过滤出j > i的点对
    mask = idxs > idx
    all_valid_distances.extend(dists[mask])

# 转为numpy数组方便后续统计
all_valid_distances = np.array(all_valid_distances)

步骤4：统计距离分布 #

用直方图分箱统计各距离的出现次数：

import matplotlib.pyplot as plt

# 设置分箱区间，比如从0到max_r分成1000个区间
bins = np.linspace(0, max_r, 1001)
counts, bin_edges = np.histogram(all_valid_distances, bins=bins)

# 绘制分布曲线
plt.figure(figsize=(10, 6))
plt.plot(bin_edges[:-1], counts, linewidth=1.5)
plt.xlabel('星系间距离')
plt.ylabel('出现次数')
plt.title('特定尺度球壳内星系距离分布')
plt.grid(axis='y', alpha=0.3)
plt.show()

关键优势 #

• 时间复杂度降至O(n log n)：KDTree的构建和查询效率远高于暴力循环。
• 内存可控：只存储符合距离范围的点对距离，避免了全量计算的内存爆炸问题。

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