Source code for titan_pylib.graph.flow.max_flow_ford_fulkerson

 1import sys
 2
 3

[docs]
 4class MaxFlowFordFulkerson:
 5
 6    def __init__(self, n: int):
 7        self.n: int = n
 8        self.m: int = 0
 9        self.G: list[list[list[int]]] = [[] for _ in range(n)]
10

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

23
24    def _dfs(self, v: int, g: int, f: int, used: list[int]) -> int:
25        if v == g:
26            return f
27        used[v] = 1
28        for i, (x, w, l) in enumerate(self.G[v]):
29            if used[x]:
30                continue
31            if w == 0:
32                continue
33            fv = self._dfs(x, g, min(f, w), used)
34            if fv > 0:
35                self.G[v][i][1] -= fv
36                self.G[x][l][1] += fv
37                return fv
38        return 0
39

[docs]
40    def max_flow(self, s: int, g: int, INF: int = 10**18) -> int:
41        """:math:`O(F|E|)"""
42        assert (
43            0 <= s < self.n
44        ), f"Indexerror: {self.__class__.__class__}.max_flow(), {s=}"
45        assert (
46            0 <= g < self.n
47        ), f"Indexerror: {self.__class__.__class__}.max_flow(), {g=}"
48        lim_rec = sys.getrecursionlimit()
49        if lim_rec < self.m:
50            sys.setrecursionlimit(self.m + 1)
51        ans = 0
52        while True:
53            used = [0] * self.n
54            f = self._dfs(s, g, INF, used)
55            if f == 0:
56                break
57            ans += f
58        return ans