如何优化Python中基于地理位置的用户匹配算法性能？

我开发了一个Python脚本，用于基于地理位置计算用户间的匹配得分。该脚本会生成用户组合、计算每组的权重平均值、权重标准差，最终为每个组合分配得分。但随着用户数量增加，脚本性能明显下降，耗时从几秒飙升到数小时。我希望获得优化该代码性能的建议，作为编程新手，若您发现其他错误或有优化建议，也欢迎告知。

原代码如下：

import random
import statistics

def comb(arr):
    if len(arr) == 2:
        yield [arr]
    for i in range(1, len(arr)):
        first = [arr[0], arr[i]]
        rest = [x for j, x in enumerate(arr) if j != 0 and j != i]
        for com in comb(rest):
            yield [first] + com

def calculate_weight(combination):
    location_1 = combination[0]['location']
    location_2 = combination[1]['location']
    weight = location_weights.get((location_1, location_2), 0)
    if weight == 0:
        weight = location_weights.get((location_2, location_1), 0)
    return weight


location_weights = {('A', 'A'): 1, ('B', 'B'): 1, ('C', 'C'): 1, ('A', 'B'): 0.3, ('B', 'C'): 0.7, ('A', 'C'): 0.5}


locations = ['A', 'B', 'C']
users = []
# Generate users with random locations
for i in range(0, 4):
    new_user = {
        'name': f'User{i}',
        'location': random.choice(locations)
    }
    users.append(new_user)

combinations = comb(users)
matches = []

for i, combination in enumerate(combinations):

    weights = []
    for pot_match in combination:
        weights.append(calculate_weight(pot_match))

    average = sum(weights) / len(weights)
    stdev = statistics.stdev(weights)

    matches.append({
        'users': combination,
        'average': average,
        'stdev': stdev,
        'score': average/(100+stdev)
    })


matches = sorted(matches, key=lambda x: x['score'], reverse=True)

for m in matches:
    print(m)

核心性能问题分析 #

原代码的comb函数生成的是所有用户的完美两两配对组合，这种组合的数量是(n-1)!!（双阶乘）：

• n=8时有105种组合
• n=10时飙升到945种
• n=20时直接超过6500万种

这种阶乘级的复杂度是性能崩溃的根本原因，小修小补无法解决大量用户场景的问题，必须更换核心实现思路。

优化建议与替代方案 #

1. 放弃全枚举，改用高效匹配算法 #

如果需求是找到最优得分的用户配对组合，不需要枚举所有可能：

• 贪心算法：每次选择当前权重最高的用户对，从剩余用户中重复该操作直到完成配对。速度极快，虽不一定全局最优，但大部分场景够用。
• 二分图最大权匹配：用匈牙利算法（Kuhn-Munkres）求解全局最优，将用户作为二分图节点、地理位置权重作为边权，Python可通过networkx库实现，效率比枚举高几个数量级。

2. 简化权重计算逻辑 #

原calculate_weight可通过排序位置对减少字典查找次数，同时精简location_weights的键：

# 先把location_weights的键统一改成排序后的元组，去掉重复项
location_weights = {('A', 'A'): 1, ('B', 'B'): 1, ('C', 'C'): 1, ('A', 'B'): 0.3, ('B', 'C'): 0.7, ('A', 'C'): 0.5}

def calculate_weight(user1, user2):
    # 排序位置对，确保不管顺序都能匹配到字典键
    loc_pair = tuple(sorted((user1['location'], user2['location'])))
    return location_weights.get(loc_pair, 0)

3. 小用户量场景的枚举优化（n≤8） #

如果必须枚举所有完美匹配，用itertools替代自定义递归函数，效率更高：

import itertools

def perfect_matches(users):
    if len(users) == 0:
        yield []
    else:
        first_user = users[0]
        # 遍历剩余用户，生成与第一个用户的配对
        for paired_user in users[1:]:
            # 筛选出未配对的剩余用户
            rest_users = [u for u in users if u != first_user and u != paired_user]
            # 递归生成剩余用户的配对组合
            for sub_match in perfect_matches(rest_users):
                yield [(first_user, paired_user)] + sub_match

4. 预计算避免重复运算 #

提前计算所有用户对的权重并缓存，后续直接读取：

# 预计算所有用户对的权重
user_pair_weights = {}
for i in range(len(users)):
    for j in range(i+1, len(users)):
        u1, u2 = users[i], users[j]
        weight = calculate_weight(u1, u2)
        user_pair_weights[(u1, u2)] = weight
        user_pair_weights[(u2, u1)] = weight

5. 其他代码细节优化 #

• 检查用户数量是否为偶数，否则处理剩余单个用户（原comb函数在奇数用户数时会报错）
• 调整得分公式中的系数：原average/(100+stdev)中100过大，标准差的影响几乎可忽略，可改为average/(1+stdev)让标准差对得分的影响更合理
• 简化测试用户生成逻辑：

  users = [{'name': f'User{i}', 'location': random.choice(locations)} for i in range(4)]
  

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