用numpy写一个匈牙利算法

好的，下面是用NumPy实现的匈牙利算法：

import numpy as np

def hungarian_algorithm(cost_matrix):
    # 初始化
    num_rows, num_cols = cost_matrix.shape
    labels = np.zeros((num_rows, num_cols))
    matching = np.full(num_cols, -1, dtype=int)
    slack = np.full(num_cols, np.inf)
    
    # 辅助函数
    def augment(col):
        visited = np.zeros(num_rows, dtype=bool)
        def dfs(row):
            visited[row] = True
            for j in range(num_cols):
                if not visited[j]:
                    reduced_cost = cost_matrix[row, j] - labels[row, col] - labels[j, j]
                    if slack[j] >= reduced_cost:
                        slack[j] = reduced_cost
                        if matching[j] == -1 or dfs(matching[j]):
                            matching[j] = col
                            return True
            return False
        return dfs
    
    # 主循环
    for i in range(num_rows):
        labels[i, :] = np.min(cost_matrix[i, :]) - cost_matrix[i, :]
        labels[:, i] = np.min(cost_matrix[:, i] - labels[:, i])
        while True:
            slack.fill(np.inf)
            visited = np.zeros(num_cols, dtype=bool)
            if augment(i)(i):
                break
            else:
                delta = np.min(slack[~visited])
                labels[visited, i] += delta
                labels[~visited, i] -= delta
                slack -= delta
    return matching

这个算法接收一个二维的NumPy数组 `cost_matrix` 作为输入，输出一个包含每列所匹配的行的索引的NumPy数组。与前面的Python实现类似，如果输出为 `[1, 0, 2]`，则第1列匹配了第1行，第2列匹配了第0行，第3列匹配了第2行。