1# from titan_pylib.graph.flow.max_flow_dinic import MaxFlowDinic
2from collections import deque
3# from titan_pylib.others.antirec import antirec
4from types import GeneratorType
5
6# ref: https://github.com/cheran-senthil/PyRival/blob/master/pyrival/misc/bootstrap.py
7# ref: https://twitter.com/onakasuita_py/status/1731535542305907041
8
9
10def antirec(func):
11 stack = []
12
13 def wrappedfunc(*args, **kwargs):
14 if stack:
15 return func(*args, **kwargs)
16 to = func(*args, **kwargs)
17 while True:
18 if isinstance(to, GeneratorType):
19 stack.append(to)
20 to = next(to)
21 else:
22 stack.pop()
23 if not stack:
24 break
25 to = stack[-1].send(to)
26 return to
27
28 return wrappedfunc
29
30
31def antirec_cache(func):
32 stack = []
33 memo = {}
34 args_list = []
35
36 def wrappedfunc(*args):
37 args_list.append(args)
38 if stack:
39 return func(*args)
40 to = func(*args)
41 while True:
42 if args_list[-1] in memo:
43 res = memo[args_list.pop()]
44 if not stack:
45 return res
46 to = stack[-1].send(res)
47 continue
48 if isinstance(to, GeneratorType):
49 stack.append(to)
50 to = next(to)
51 else:
52 memo[args_list.pop()] = to
53 stack.pop()
54 if not stack:
55 break
56 to = stack[-1].send(to)
57 return to
58
59 return wrappedfunc
60
61
62class MaxFlowDinic:
63 """mf.G[v]:= [x, cap, ind, flow]"""
64
65 def __init__(self, n: int):
66 self.n: int = n
67 self.G: list[list[list[int]]] = [[] for _ in range(n)]
68 self.level = [-1] * n
69
70 def add_edge(self, u: int, v: int, w: int) -> None:
71 assert (
72 0 <= u < self.n
73 ), f"Indexerror: {self.__class__.__name__}.add_edge({u}, {v})"
74 assert (
75 0 <= v < self.n
76 ), f"Indexerror: {self.__class__.__name__}.add_edge({u}, {v})"
77 G_u = len(self.G[u])
78 G_v = len(self.G[v])
79 self.G[u].append([v, w, G_v, 0])
80 self.G[v].append([u, 0, G_u, 0])
81
82 def _bfs(self, s: int) -> None:
83 level = self.level
84 for i in range(len(level)):
85 level[i] = -1
86 dq = deque([s])
87 level[s] = 0
88 while dq:
89 v = dq.popleft()
90 for x, w, _, _ in self.G[v]:
91 if w > 0 and level[x] == -1:
92 level[x] = level[v] + 1
93 dq.append(x)
94 self.level = level
95
96 @antirec
97 def _dfs(self, v: int, g: int, f: int):
98 if v == g:
99 yield f
100 else:
101 for i in range(self.it[v], len(self.G[v])):
102 self.it[v] += 1
103 x, w, rev, _ = self.G[v][i]
104 if w > 0 and self.level[v] < self.level[x]:
105 fv = yield self._dfs(x, g, min(f, w))
106 if fv > 0:
107 self.G[v][i][3] += f
108 self.G[x][rev][3] -= f
109 self.G[v][i][1] -= fv
110 self.G[x][rev][1] += fv
111 yield fv
112 break
113 else:
114 yield 0
115
116 def max_flow(self, s: int, g: int, INF: int = 10**18) -> int:
117 """:math:`O(V^2 E)`"""
118 assert (
119 0 <= s < self.n
120 ), f"Indexerror: {self.__class__.__class__}.max_flow(), {s=}"
121 assert (
122 0 <= g < self.n
123 ), f"Indexerror: {self.__class__.__class__}.max_flow(), {g=}"
124 ans = 0
125 while True:
126 self._bfs(s)
127 if self.level[g] < 0:
128 break
129 self.it = [0] * self.n
130 while True:
131 f = self._dfs(s, g, INF)
132 if f == 0:
133 break
134 ans += f
135 return ans