Python实现无放回K近邻最优匹配：为项目地块分配对照地块

问题背景 #

需要为每个项目地块（PP）分配2个对照地块，规则如下：

1. 第一轮分配：为每个PP选择距离最近的对照地块，若该对照已被本轮其他PP分配，则选择下一个最近的。
2. 第二轮分配：待所有PP完成首次分配后，重复第一轮的规则，为每个PP分配第二个对照地块。

输入为包含control_plot_id及各PP距离的DataFrame，期望输出包含control_plot_id、PP、dist三列的结果。

解决方案代码 #

假设输入数据为宽格式（control_plot_id为一列，其余列对应不同PP，值为距离），以下是实现逻辑的Python代码：

import pandas as pd

# 读取原始数据（替换为你的数据路径/数据源）
df = pd.read_csv("control_plot_distances.csv")

# 将宽格式数据转换为长格式，便于按PP和距离排序
long_format_df = df.melt(
    id_vars="control_plot_id",
    var_name="PP",
    value_name="dist"
)

# 对每个PP的对照地块按距离从小到大排序
sorted_candidates = long_format_df.groupby("PP").apply(
    lambda x: x.sort_values("dist", ascending=True)
).reset_index(drop=True)

# 初始化结果列表
final_assignments = []

# --------------------------
# 第一轮分配：第一个对照地块
# --------------------------
assigned_round1 = set()
for pp in sorted_candidates["PP"].unique():
    pp_options = sorted_candidates[sorted_candidates["PP"] == pp]
    for _, row in pp_options.iterrows():
        control_id = row["control_plot_id"]
        if control_id not in assigned_round1:
            final_assignments.append({
                "control_plot_id": control_id,
                "PP": pp,
                "dist": row["dist"]
            })
            assigned_round1.add(control_id)
            break

# --------------------------
# 第二轮分配：第二个对照地块
# --------------------------
assigned_round2 = set()
for pp in sorted_candidates["PP"].unique():
    pp_options = sorted_candidates[sorted_candidates["PP"] == pp]
    for _, row in pp_options.iterrows():
        control_id = row["control_plot_id"]
        if control_id not in assigned_round2:
            final_assignments.append({
                "control_plot_id": control_id,
                "PP": pp,
                "dist": row["dist"]
            })
            assigned_round2.add(control_id)
            break

# 转换为最终结果DataFrame
result_df = pd.DataFrame(final_assignments)
print(result_df)

代码逻辑说明 #

1. 格式转换：将宽格式数据转为长格式，把每个PP的距离信息整合为单条记录，方便后续排序和遍历。
2. 排序候选：对每个PP的所有对照地块按距离升序排序，确保优先选择最近的对照。
3. 轮次分配：
   • 每轮维护一个集合记录已被分配的对照ID，避免同一轮中重复分配给不同PP。
   • 遍历每个PP的候选列表，找到第一个未被当前轮次分配的对照，加入结果集并标记为已分配。
4. 结果输出：合并两轮分配结果，得到每个PP对应的两个对照地块记录。

边界情况处理 #

如果存在PP的候选对照数量不足（比如对照地块总数少于PP数量），可以在遍历候选时添加判断，避免无终止循环：

# 在轮次分配的循环中添加判断
found = False
for _, row in pp_options.iterrows():
    control_id = row["control_plot_id"]
    if control_id not in assigned_round1:
        final_assignments.append({
            "control_plot_id": control_id,
            "PP": pp,
            "dist": row["dist"]
        })
        assigned_round1.add(control_id)
        found = True
        break
if not found:
    print(f"警告：PP {pp} 无可用对照地块")

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