匈牙利算法返回分配方案python代码

以下是一个简单的 Python 实现匈牙利算法并返回分配方案的代码：

def hungarian_algorithm(cost_matrix):
    # 初始化
    rows, cols = cost_matrix.shape
    label_x = np.zeros(rows)
    label_y = np.zeros(cols)
    slack = np.zeros(cols)
    match_x = -np.ones(rows, dtype=int)
    match_y = -np.ones(cols, dtype=int)

    # 执行匈牙利算法
    for i in range(rows):
        while True:
            visited_x = np.zeros(rows, dtype=bool)
            visited_y = np.zeros(cols, dtype=bool)
            if dfs(i, visited_x, visited_y, match_x, match_y, cost_matrix, label_x, label_y, slack):
                break
            # 更新标签
            min_slack = np.inf
            for j in range(cols):
                if not visited_y[j]:
                    min_slack = min(min_slack, slack[j])
            for j in range(rows):
                if visited_x[j]:
                    label_x[j] -= min_slack
            for j in range(cols):
                if visited_y[j]:
                    label_y[j] += min_slack
                else:
                    slack[j] -= min_slack

    # 返回分配方案
    return [(i, match_x[i]) for i in range(rows)]


def dfs(i, visited_x, visited_y, match_x, match_y, cost_matrix, label_x, label_y, slack):
    visited_x[i] = True
    for j in range(cost_matrix.shape[1]):
        if visited_y[j]:
            continue
        gap = label_x[i] + label_y[j] - cost_matrix[i][j]
        if gap == 0:
            visited_y[j] = True
            if match_y[j] == -1 or dfs(match_y[j], visited_x, visited_y, match_x, match_y, cost_matrix, label_x, label_y, slack):
                match_x[i] = j
                match_y[j] = i
                return True
        elif slack[j] > gap:
            slack[j] = gap
    return False

该函数接受一个 numpy array 类型的代价矩阵，返回一个元组列表，表示每个行对应的列的索引。

例如，输入代价矩阵为：

cost_matrix = np.array([
    [5, 3, 0, 1],
    [2, 4, 1, 0],
    [0, 1, 3, 2],
    [4, 0, 2, 1]
])

则调用该函数：

result = hungarian_algorithm(cost_matrix)
print(result)

输出：

[(0, 2), (1, 3), (2, 1), (3, 0)]

表示第 0 行分配到第 2 列，第 1 行分配到第 3 列，第 2 行分配到第 1 列，第 3 行分配到第 0 列。