代码

###python3

INF = int(1e18)

class Solution:
    def minimumTotalDistance(self, robot: List[int], factory: List[List[int]]) -> int:
        n, m = len(robot), len(factory)
        STRAT, END = n + m + 3, n + m + 4
        mcmf = MinCostMaxFlow(n + m + 10, STRAT, END)
        for i in range(n):
            mcmf.addEdge(STRAT, i, 1, 0)
        for i in range(n):
            for j in range(m):
                mcmf.addEdge(i, n + j, 1, abs(robot[i] - factory[j][0]))
        for i in range(m):
            mcmf.addEdge(n + i, END, factory[i][1], 0)
        return mcmf.work()[1]



class Edge:
    __slots__ = ("fromV", "toV", "cap", "cost", "flow")

    def __init__(self, fromV: int, toV: int, cap: int, cost: int, flow: int) -> None:
        self.fromV = fromV
        self.toV = toV
        self.cap = cap
        self.cost = cost
        self.flow = flow


class MinCostMaxFlow:
    """最小费用流的连续最短路算法复杂度为流量*最短路算法复杂度"""

    __slots__ = ("_n", "_start", "_end", "_edges", "_reGraph", "_dist", "_visited", "_curEdges")

    def __init__(self, n: int, start: int, end: int):
        """
        Args:
            n (int): 包含虚拟点在内的总点数
            start (int): (虚拟)源点
            end (int): (虚拟)汇点
        """
        assert 0 <= start < n and 0 <= end < n
        self._n = n
        self._start = start
        self._end = end
        self._edges: List["Edge"] = []
        self._reGraph: List[List[int]] = [[] for _ in range(n + 10)]  # 残量图存储的是边的下标

        self._dist = [INF] * (n + 10)
        self._visited = [False] * (n + 10)
        self._curEdges = [0] * (n + 10)

    def addEdge(self, fromV: int, toV: int, cap: int, cost: int) -> None:
        """原边索引为i 反向边索引为i^1"""
        self._edges.append(Edge(fromV, toV, cap, cost, 0))
        self._edges.append(Edge(toV, fromV, 0, -cost, 0))
        len_ = len(self._edges)
        self._reGraph[fromV].append(len_ - 2)
        self._reGraph[toV].append(len_ - 1)

    def work(self) -> Tuple[int, int]:
        """
        Returns:
            Tuple[int, int]: [最大流,最小费用]
        """
        maxFlow, minCost = 0, 0
        while self._spfa():
            # !如果流量限定为1，那么一次dfs只会找到一条费用最小的增广流
            # !如果流量限定为INF，那么一次dfs不只会找到一条费用最小的增广流
            flow = self._dfs(self._start, self._end, INF)
            maxFlow += flow
            minCost += flow * self._dist[self._end]
        return maxFlow, minCost

    def slope(self) -> List[Tuple[int, int]]:
        """
        Returns:
            List[Tuple[int, int]]: 流量为a时,最小费用是b
        """
        res = [(0, 0)]
        flow, cost = 0, 0
        while self._spfa():
            deltaFlow = self._dfs(self._start, self._end, INF)
            flow += deltaFlow
            cost += deltaFlow * self._dist[self._end]
            res.append((flow, cost))  # type: ignore
        return res

    def _spfa(self) -> bool:
        """spfa沿着最短路寻找增广路径  有负cost的边不能用dijkstra"""
        n, start, end, edges, reGraph, visited = (
            self._n,
            self._start,
            self._end,
            self._edges,
            self._reGraph,
            self._visited,
        )

        self._curEdges = [0] * n
        self._dist = dist = [INF] * n
        dist[start] = 0
        queue = deque([start])

        while queue:
            cur = queue.popleft()
            visited[cur] = False
            for edgeIndex in reGraph[cur]:
                edge = edges[edgeIndex]
                cost, remain, next = edge.cost, edge.cap - edge.flow, edge.toV
                if remain > 0 and dist[cur] + cost < dist[next]:
                    dist[next] = dist[cur] + cost
                    if not visited[next]:
                        visited[next] = True
                        if queue and dist[queue[0]] > dist[next]:
                            queue.appendleft(next)
                        else:
                            queue.append(next)

        return dist[end] != INF

    def _dfs(self, cur: int, end: int, flow: int) -> int:
        if cur == end:
            return flow

        visited, reGraph, curEdges, edges, dist = (
            self._visited,
            self._reGraph,
            self._curEdges,
            self._edges,
            self._dist,
        )

        visited[cur] = True
        res = flow
        index = curEdges[cur]
        while res and index < len(reGraph[cur]):
            edgeIndex = reGraph[cur][index]
            next, remain = edges[edgeIndex].toV, edges[edgeIndex].cap - edges[edgeIndex].flow
            if remain > 0 and not visited[next] and dist[next] == dist[cur] + edges[edgeIndex].cost:
                delta = self._dfs(next, end, remain if remain < res else res)
                res -= delta
                edges[edgeIndex].flow += delta
                edges[edgeIndex ^ 1].flow -= delta
            curEdges[cur] += 1
            index = curEdges[cur]

        visited[cur] = False
        return flow - res

###python3

INF = int(1e18)

class Solution:
    def minimumTotalDistance(self, robot: List[int], factory: List[List[int]]) -> int:
        boys, girls = robot, []  
        for pos, limit in factory:
            girls.extend([pos] * limit)  
        costMatrix = [[0] * len(girls) for _ in range(len(boys))]
        for i, pos1 in enumerate(boys):
            for j, pos2 in enumerate(girls):
                costMatrix[i][j] = -abs(pos1 - pos2)  # 最大权匹配转换为最小权匹配
        return -KM(costMatrix)[0]

        
def KM(costMatrix: List[List[int]]) -> Tuple[int, Tuple[List[int], List[int]]]:
    """KM算法求带权二分图的最大权匹配

    Args
    ----------
    costMatrix (List[List[int]]):
        二分图的权值矩阵,不存在的边应初始化为`-INF`

    Returns
    ----------
    Tuple[int, Tuple[List[int], List[int]]]:
        `最大权匹配值, 匹配对的行索引、列索引`

    Examples
    ----------
    >>> costMatrix = [[1, 2, 3], [2, 4, 6], [3, 6, 9]]
    >>> maxSum, (rows, cols) = KuhnMunkres(costMatrix)
    >>> maxSum
    14
    >>> rows cols
    [0, 1, 2] [0, 1, 2]
    >>> sum(costMatrix[i][j] for i, j in zip(rows, cols))
    14
    """
    max_ = max(len(costMatrix), len(costMatrix[0]))
    _match = [-1] * max_  # 记录每个女生匹配到的男生 如果没有则为-1
    _graph = costMatrix  # 记录每个男生和每个女生之间的`好感度`
    _visitedBoy = set()  # 记录每一轮匹配匹配过的男生
    _visitedGirl = set()  # 记录每一轮匹配匹配过的女生
    _expBoy = [max(row) for row in costMatrix]  # 每个男生的期望值
    _expGirl = [0] * max_  # 每个女生的期望值，为0表示只要有一个男生就可以
    _slack = []  # 记录每个女生如果能被男生倾心最少还需要多少期望值
    _pre = []
    _row = len(costMatrix)
    _col = len(costMatrix[0])

    def dfs(boy: int) -> int:
        _visitedBoy.add(boy)
        for girl in range(_col):
            if girl in _visitedGirl:
                continue
            delta = _expBoy[boy] + _expGirl[girl] - _graph[boy][girl]
            # 符合要求
            if delta == 0:
                _visitedGirl.add(girl)
                _pre[girl + _row] = boy
                if _match[girl] == -1:
                    return girl + _row
                _pre[_match[girl]] = girl + _row
                nextRes = dfs(_match[girl])  # 找到增广
                if nextRes > 0:
                    return nextRes
            # 女生要得到男生的倾心 还需多少期望值
            elif _slack[boy] > delta:
                _slack[boy] = delta

        return -1

    for boy in range(_row):
        _visitedBoy.clear()
        _visitedGirl.clear()
        _slack = [INF] * _col
        _pre = [-1] * (_row + _col)
        visited = False
        cand = -1
        # 记录每轮匹配中男生女生是否被尝试匹配过
        while True:
            if not visited:
                cand = dfs(boy)
                visited = True
            else:
                for r in range(_row):
                    if _slack[r] == 0:
                        _slack[r] = INF
                        cand = dfs(r)
                        if cand > 0:
                            break

            if cand > 0:
                tmp = cand
                while tmp > 0:
                    _match[tmp - _row] = _pre[tmp]
                    tmp = _pre[_pre[tmp]]
                break
            else:
                # 如果不能找到 就降低期望值
                # 最小可降低的期望值
                delta = INF
                for c in range(_row):
                    if c in _visitedBoy and _slack[c] < delta:
                        delta = _slack[c]
                for r in range(_row):
                    if r in _visitedBoy:
                        # 所有访问过的男生降低期望值
                        _expBoy[r] -= delta
                        _slack[r] -= delta
                for c in range(_col):
                    if c in _visitedGirl:
                        # 所有访问过的女生增加期望值
                        _expGirl[c] += delta

    # 匹配完成 求出所有配对的好感度的和
    res, rows, cols = 0, [], []
    for girl, boy in enumerate(_match):
        if boy != -1:
            res += _graph[boy][girl]
            rows.append(boy)
            cols.append(girl)
    return res, (rows, cols)