Source code for titan_pylib.graph.flow.max_flow_dinic

 1from collections import deque
 2from titan_pylib.others.antirec import antirec
 3
 4

[docs]
 5class MaxFlowDinic:
 6    """mf.G[v]:= [x, cap, ind, flow]"""
 7
 8    def __init__(self, n: int):
 9        self.n: int = n
10        self.G: list[list[list[int]]] = [[] for _ in range(n)]
11        self.level = [-1] * n
12

[docs]
13    def add_edge(self, u: int, v: int, w: int) -> None:
14        assert (
15            0 <= u < self.n
16        ), f"Indexerror: {self.__class__.__name__}.add_edge({u}, {v})"
17        assert (
18            0 <= v < self.n
19        ), f"Indexerror: {self.__class__.__name__}.add_edge({u}, {v})"
20        G_u = len(self.G[u])
21        G_v = len(self.G[v])
22        self.G[u].append([v, w, G_v, 0])
23        self.G[v].append([u, 0, G_u, 0])

24
25    def _bfs(self, s: int) -> None:
26        level = self.level
27        for i in range(len(level)):
28            level[i] = -1
29        dq = deque([s])
30        level[s] = 0
31        while dq:
32            v = dq.popleft()
33            for x, w, _, _ in self.G[v]:
34                if w > 0 and level[x] == -1:
35                    level[x] = level[v] + 1
36                    dq.append(x)
37        self.level = level
38
39    @antirec
40    def _dfs(self, v: int, g: int, f: int):
41        if v == g:
42            yield f
43        else:
44            for i in range(self.it[v], len(self.G[v])):
45                self.it[v] += 1
46                x, w, rev, _ = self.G[v][i]
47                if w > 0 and self.level[v] < self.level[x]:
48                    fv = yield self._dfs(x, g, min(f, w))
49                    if fv > 0:
50                        self.G[v][i][3] += f
51                        self.G[x][rev][3] -= f
52                        self.G[v][i][1] -= fv
53                        self.G[x][rev][1] += fv
54                        yield fv
55                        break
56            else:
57                yield 0
58

[docs]
59    def max_flow(self, s: int, g: int, INF: int = 10**18) -> int:
60        """:math:`O(V^2 E)`"""
61        assert (
62            0 <= s < self.n
63        ), f"Indexerror: {self.__class__.__class__}.max_flow(), {s=}"
64        assert (
65            0 <= g < self.n
66        ), f"Indexerror: {self.__class__.__class__}.max_flow(), {g=}"
67        ans = 0
68        while True:
69            self._bfs(s)
70            if self.level[g] < 0:
71                break
72            self.it = [0] * self.n
73            while True:
74                f = self._dfs(s, g, INF)
75                if f == 0:
76                    break
77                ans += f
78        return ans