匈牙利算法

用于寻找二分图的最大匹配

import numpy as np

class HungarianAlgorithm:
    def __init__(self, num1, num2, links):
        self.num1 = num1  # number of part 1
        self.num2 = num2  # number of part 2
        self.adjacent_matrix = np.zeros((num1, num2), dtype=int)
        for m, n in links:
            self.adjacent_matrix[m, n] = 1
        self.visited = np.zeros((num2,), dtype=bool)  # for part 2
        self.lut_2to1 = np.full((num2,), -1, dtype=int)  # from idx2 to idx1

    def find_max_matching(self):
        count = 0
        for idx1 in range(self.num1):
            self.visited[...] = False
            if self.match(idx1):
                count += 1
        return count

    def match(self, idx1):
        for idx2 in range(self.num2):
            if self.adjacent_matrix[idx1, idx2] and not self.visited[idx2]:
                self.visited[idx2] = True
                if self.lut_2to1[idx2] == -1 or self.match(self.lut_2to1[idx2]):
                    self.lut_2to1[idx2] = idx1
                    return True
        return False

if __name__ == '__main__':
    num1 = 4
    num2 = 5
    links = [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 1], [2, 2], [2, 3], [3, 0]]
    alg = HungarianAlgorithm(num1, num2, links)
    count = alg.find_max_matching()