| 2 | |
| 3 | |
| 4 | class Dinic: |
| 5 | def __init__(self, n): |
| 6 | self.lvl = [0] * n |
| 7 | self.ptr = [0] * n |
| 8 | self.q = [0] * n |
| 9 | self.adj = [[] for _ in range(n)] |
| 10 | |
| 11 | """ |
| 12 | Here we will add our edges containing with the following parameters: |
| 13 | vertex closest to source, vertex closest to sink and flow capacity |
| 14 | through that edge ... |
| 15 | """ |
| 16 | |
| 17 | def add_edge(self, a, b, c, rcap=0): |
| 18 | self.adj[a].append([b, len(self.adj[b]), c, 0]) |
| 19 | self.adj[b].append([a, len(self.adj[a]) - 1, rcap, 0]) |
| 20 | |
| 21 | # This is a sample depth first search to be used at max_flow |
| 22 | def depth_first_search(self, vertex, sink, flow): |
| 23 | if vertex == sink or not flow: |
| 24 | return flow |
| 25 | |
| 26 | for i in range(self.ptr[vertex], len(self.adj[vertex])): |
| 27 | e = self.adj[vertex][i] |
| 28 | if self.lvl[e[0]] == self.lvl[vertex] + 1: |
| 29 | p = self.depth_first_search(e[0], sink, min(flow, e[2] - e[3])) |
| 30 | if p: |
| 31 | self.adj[vertex][i][3] += p |
| 32 | self.adj[e[0]][e[1]][3] -= p |
| 33 | return p |
| 34 | self.ptr[vertex] = self.ptr[vertex] + 1 |
| 35 | return 0 |
| 36 | |
| 37 | # Here we calculate the flow that reaches the sink |
| 38 | def max_flow(self, source, sink): |
| 39 | flow, self.q[0] = 0, source |
| 40 | for l in range(31): # l = 30 maybe faster for random data # noqa: E741 |
| 41 | while True: |
| 42 | self.lvl, self.ptr = [0] * len(self.q), [0] * len(self.q) |
| 43 | qi, qe, self.lvl[source] = 0, 1, 1 |
| 44 | while qi < qe and not self.lvl[sink]: |
| 45 | v = self.q[qi] |
| 46 | qi += 1 |
| 47 | for e in self.adj[v]: |
| 48 | if not self.lvl[e[0]] and (e[2] - e[3]) >> (30 - l): |
| 49 | self.q[qe] = e[0] |
| 50 | qe += 1 |
| 51 | self.lvl[e[0]] = self.lvl[v] + 1 |
| 52 | |
| 53 | p = self.depth_first_search(source, sink, INF) |
| 54 | while p: |
| 55 | flow += p |
| 56 | p = self.depth_first_search(source, sink, INF) |
| 57 | |
| 58 | if not self.lvl[sink]: |
| 59 | break |
| 60 | |
| 61 | return flow |