Python中使用最大权闭合子图算法解决选课问题

from collections import defaultdict, deque

class Graph:
    def __init__(self, n):
        self.n = n
        self.graph = defaultdict(list)
        self.weights = [0] * (n+1)

    def add_edge(self, u, v):
        self.graph[u].append(v)

    def add_weight(self, u, w):
        self.weights[u] = w

    def max_closure(self):
        # 在原图基础上添加超级源点和超级汇点
        s = self.n + 1
        t = self.n + 2
        for i in range(1, self.n+1):
            if self.weights[i] > 0:
                # 如果节点权值大于0，表示这是一门可以选的课程
                self.add_edge(s, i)
            elif self.weights[i] < 0:
                # 如果节点权值小于0，表示这是一门必须选的课程
                self.add_edge(i, t)

        # 对于任意一门课程 i，如果 j 是其先修课程，就从 i 指向 j 连一条有向边
        for i in range(1, self.n+1):
            for j in self.graph[i]:
                self.add_edge(i, j)

        # 对新建的超级源点和超级汇点分别计算其余图中每个节点与其的连边容量
        capacities = defaultdict(int)
        for i in range(1, self.n+1):
            if self.weights[i] > 0:
                # 如果节点权值大于0，表示这是一门可以选的课程
                capacities[s, i] = self.weights[i]
            elif self.weights[i] < 0:
                # 如果节点权值小于0，表示这是一门必须选的课程
                capacities[i, t] = -self.weights[i]
            else:
                # 如果节点权值为0，表示这是一门普通的课程
                capacities[i, i] = float('inf')

        # 使用 BFS 和 DFS 处理残留网络
        def bfs(s, t, parent):
            visited = set()
            queue = deque([s])
            visited.add(s)
            while queue:
                u = queue.popleft()
                for v in self.graph[u]:
                    if (u, v) not in capacities:
                        continue
                    if v not in visited:
                        queue.append(v)
                        visited.add(v)
                        parent[v] = u
                        if v == t:
                            return True
            return False

        def dfs(u, t, parent, visited):
            if u not in visited:
                visited.add(u)
                if u == t:
                    return True
                for v in self.graph[u]:
                    if (u, v) in capacities and dfs(v, t, parent, visited):
                        parent[v] = u
                        return True
            return False

        # 计算最大流量
        parent = {}
        max_flow = 0
        while bfs(s, t, parent):
            path_flow = float('inf')
            s_ = t
            while s_ != s:
                path_flow = min(path_flow, capacities[parent[s_], s_])
                s_ = parent[s_]
            max_flow += path_flow
            t_ = t
            while t_ != s:
                u_ = parent[t_]
                capacities[u_, t_] -= path_flow
                capacities[t_, u_] += path_flow
                t_ = u_
            parent = {}

        # 根据连边容量判断最小割，并输出选出的课程
        output = []
        for i, j in capacities:
            if capacities[i, j] == 0 and i != s and j != t:
                output.append(j)
        return sorted(output)

g = Graph(12)
g.add_edge(1, 2)
g.add_edge(2, 3)
g.add_edge(2, 4)
g.add_edge(3, 5)
g.add_edge(4, 5)
g.add_edge(5, 6)
g.add_edge(7, 8)
g.add_edge(8, 9)
g.add_edge(8, 10)
g.add_edge(9, 11)
g.add_edge(10, 11)
g.add_edge(11, 12)
g.add_weight(1, 2)
g.add_weight(2, 1)
g.add_weight(3, 1)
g.add_weight(4, 1)
g.add_weight(5, -5)
g.add_weight(6, 2)
g.add_weight(7, 2)
g.add_weight(8, 1)
g.add_weight(9, 1)
g.add_weight(10, 1)
g.add_weight(11, -3)
g.add_weight(12, 2)

print(g.max_closure())