咨询：如何计算三组处理组中最小的组内Trio距离？

DevOps 小满

2026-8-13

问题描述

我有如下结构的Pandas DataFrame：

Output  var1        var2        var3
1   0.487981    0.297929    0.214090    
1   0.945660    0.031666    0.022674
2   0.119845    0.828661    0.051495
2   0.095186    0.852232    0.052582
3   0.059520    0.053307    0.887173
3   0.091049    0.342226    0.566725
3   0.119295    0.414376    0.466329
... ... ... ... ...

其中：

Output是处理组标记（取值1、2、3）
var1/var2/var3是三个倾向得分，我将其视为三维空间中的x/y/z坐标，每个样本对应空间中的一个点

我的需求是：计算跨处理组的Trio组合（每个组合必须包含1、2、3组各一个样本）的组内Trio距离，找到距离最小的Trio。这个方法来自Rassen等人的论文《Matching by Propensity Score in Cohort Studies with Three Treatment Groups》，论文提到该距离类似三角形周长，但我不确定具体定义。

我找到过相关的Java实现，但需要用Python重写。目前已经有计算两点欧氏距离的代码：

import numpy as np
import itertools

# 计算两点间欧氏距离
distance = lambda p1, p2: np.sqrt(np.sum((p1 - p2) ** 2))

def min_distance(cloud):
    pairs = itertools.combinations(cloud, 2)
    return np.min(map(lambda pair: distance(*pair), pairs))

但不确定三点的Trio距离是否需要计算所有两两距离之和，特此咨询正确的实现方式。

实现方案

1. 明确Trio距离的定义

根据你提到的论文和对应Java实现的逻辑，Trio距离的计算方式确实是三个点两两之间的欧氏距离之和（即三角形的周长）。因为是三处理组匹配，每个有效Trio必须包含1、2、3组各一个样本，不能是同组内的三点组合。

2. Python实现步骤

步骤1：分组提取各组样本点

先将DataFrame按Output分组，提取每个组的三维坐标数组：

import pandas as pd
import numpy as np
import itertools

# 假设你的DataFrame名为df
df = pd.read_csv("你的数据文件路径.csv")

# 按处理组分组，提取坐标
group_1 = df[df["Output"] == 1][["var1", "var2", "var3"]].values
group_2 = df[df["Output"] == 2][["var1", "var2", "var3"]].values
group_3 = df[df["Output"] == 3][["var1", "var2", "var3"]].values

步骤2：定义Trio距离计算函数

基于两点距离函数，实现Trio距离的计算：

def euclidean_distance(p1, p2):
    return np.sqrt(np.sum((p1 - p2) ** 2))

def trio_distance(p1, p2, p3):
    # 计算三个两两距离的和
    d12 = euclidean_distance(p1, p2)
    d13 = euclidean_distance(p1, p3)
    d23 = euclidean_distance(p2, p3)
    return d12 + d13 + d23

步骤3：遍历所有跨组Trio，找到最小距离的组合

生成所有可能的跨组三元组（每组取一个样本），计算每个Trio的距离，记录最小距离对应的组合：

# 生成所有跨组三元组
all_triples = itertools.product(group_1, group_2, group_3)

min_total_distance = float('inf')
best_triple = None

for triple in all_triples:
    p1, p2, p3 = triple
    dist = trio_distance(p1, p2, p3)
    if dist < min_total_distance:
        min_total_distance = dist
        best_triple = (p1, p2, p3)

# 输出结果
print(f"最小组内Trio距离: {min_total_distance:.6f}")
print("对应的三个样本坐标:")
print(f"组1样本: {best_triple[0]}")
print(f"组2样本: {best_triple[1]}")
print(f"组3样本: {best_triple[2]}")

3. 性能优化提示

如果各组样本量较大（比如每组超过100个样本），直接遍历所有组合会非常耗时（比如100100100=1e6次计算）。可以考虑以下优化：

先对每组样本进行降维（比如PCA），减少计算量
使用KD-Tree等空间索引结构，快速查找近似最近邻，缩小候选范围
并行计算：用multiprocessing模块加速遍历过程

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