Problem p02561 (python)

Fixed submission:
from collections import deque


class MF_graph(object):
    def __init__(self, n):
        self.n = n
        self.g = [[] for _ in range(n)]  # to, rev, cap
        self.pos = []

    def add_edge(self, frm, to, cap):
        m = len(self.pos)
        self.pos.append((frm, len(self.g[frm])))
        self.g[frm].append([to, len(self.g[to]), cap])
        self.g[to].append([frm, len(self.g[frm]) - 1, 0])
        return m

    def get_edge(self, i):
        e_to, e_rev, e_cap = self.g[self.pos[i][0]][self.pos[i][1]]
        re_to, _, re_cap = self.g[e_to][e_rev]
        # from, to, cap, flow
        return (re_to, e_to, e_cap + re_cap, re_cap)

    def edges(self):
        m = len(self.pos)
        for i in range(m):
            yield self.get_edge(i)

    def change_edge(self, i, new_cap, new_flow):
        f, s = self.pos[i]
        rf, rs, _ = self.g[f][s]
        self.g[f][s][2] = new_cap - new_flow
        self.g[rf][rs][2] = new_flow
        return

    def dfs(self, s, v, up):
        if v == s:
            return up
        res = 0
        level_v = self.level[v]
        for i in range(self.iter[v], len(self.g[v])):
            u_to, u_rev, _ = self.g[v][i]
            if level_v <= self.level[u_to] or self.g[u_to][u_rev][2] == 0:
                continue
            d = self.dfs(s, u_to, min(up - res, self.g[u_to][u_rev][2]))
            if d <= 0:
                continue
            self.g[v][i][2] += d
            self.g[u_to][u_rev][2] -= d
            res += d
            if res == up:
                break
        return res

    def flow(self, s, t, flow_limit=10**18):
        self.iter = [0] * self.n

        flow = 0
        while flow < flow_limit:
            self.level = [-1] * self.n
            self.level[s] = 0
            que = deque([s])
            while que:
                v = que.popleft()
                for u_to, _, u_cap in self.g[v]:
                    if u_cap == 0 or self.level[u_to] >= 0:
                        continue
                    self.level[u_to] = self.level[v] + 1
                    if u_to == t:
                        break
                    que.append(u_to)

            if self.level[t] == -1:
                break
            self.iter = [0] * self.n
            while flow < flow_limit:
                f = self.dfs(s, t, flow_limit - flow)
                if not f:
                    break
                flow += f
        return flow

    def min_cut(self, s):
        visited = [False] * self.n
        que = deque([s])
        while que:
            v = que.popleft()
            visited[v] = True
            for u_to, _, u_cap in self.g[v]:
                if u_cap and (not visited[u_to]):
                    visited[u_to] = True
                    que.append(u_to)
        return visited


if __name__ == "__main__":
    import sys
    input = sys.stdin.readline
    sys.setrecursionlimit(10 ** 7)

    N, M = map(int, input().split())
    grid = [list(input().rstrip()) for _ in range(N)]

    g = MF_graph(N * M + 2)
    source = N * M
    sink = N * M + 1

    for i in range(N):
        for j in range(M):
            if grid[i][j] == "#":
                continue
            v = i * M + j
            if (i + j) % 2 == 0:
                g.add_edge(source, v, 1)
            else:
                g.add_edge(v, sink, 1)

    for i in range(N):
        for j in range(M):
            if (i + j) % 2 or grid[i][j] == "#":
                continue
            v0 = i * M + j
            if i and grid[i - 1][j] == ".":
                v1 = (i - 1) * M + j
                g.add_edge(v0, v1, 1)
            if j and grid[i][j - 1] == ".":
                v1 = i * M + j - 1
                g.add_edge(v0, v1, 1)
            if i + 1 < N and grid[i + 1][j] == ".":
                v1 = (i + 1) * M + j
                g.add_edge(v0, v1, 1)
            if j + 1 < M and grid[i][j + 1] == ".":
                v1 = i * M + j + 1
                g.add_edge(v0, v1, 1)

    res = g.flow(source, sink)

    for frm, to, cap, flow in g.edges():
        if frm == source or to == sink or flow == 0:
            continue
        i0, j0 = divmod(frm, M)
        i1, j1 = divmod(to, M)
        if i0 == i1 + 1:
            grid[i1][j1] = "v"
            grid[i0][j0] = "^"
        elif j0 == j1 + 1:
            grid[i1][j1] = ">"
            grid[i0][j0] = "<"
        elif i0 == i1 - 1:
            grid[i1][j1] = "^"
            grid[i0][j0] = "v"
        else:
            grid[i1][j1] = "<"
            grid[i0][j0] = ">"

    print(res)
    for g in grid:
        print("".join(g))

Buggy submission (verdict: Runtime Error):
class BipartiteMatching:
    def __init__(self, N, M):
        self.N = N
        self.M = M
        self.pairA = [-1] * N
        self.pairB = [-1] * M
        self.edge = [[] for _ in range(N)]
        self.was = [0] * N
        self.iter = 0

    def add(self, s, t):
        self.edge[s].append(t)

    def _DFS(self, s):
        self.was[s] = self.iter
        for t in self.edge[s]:
            if self.pairB[t] == -1:
                self.pairA[s] = t
                self.pairB[t] = s
                return True
        for t in self.edge[s]:
            if self.was[self.pairB[t]] != self.iter and self._DFS(self.pairB[t]):
                self.pairA[s] = t
                self.pairB[t] = s
                return True
        return False

    def solve(self):
        res = 0
        while True:
            self.iter += 1
            found = 0
            for i in range(self.N):
                if self.pairA[i] == -1 and self._DFS(i):
                    found += 1
            if not found:
                break
            res += found
        return res


DX = (-1, 0, 1, 0, -1, -1, 1, 1)
DY = (0, 1, 0, -1, -1, 1, -1, 1)
DX = DX[:4]
DY = DY[:4]


H, W = map(int, input().split())
grid = [list(input()) for _ in range(H)]
U = H * W


def num(x, y):
    return W * x + y


def coo(n):
    return divmod(n, W)


match = BipartiteMatching(U, U)
for i, si in enumerate(grid):
    for j, s in enumerate(si):
        if s == "#":
            continue
        if (i + j) % 2 == 1:
            continue

        x = num(i, j)
        for dx, dy in zip(DX, DY):
            nx = i + dx
            ny = j + dy
            if nx < 0 or nx >= H or ny < 0 or ny >= W:
                continue
            if grid[nx][ny] == ".":
                y = num(nx, ny)
                match.add(x, y)

ans = match.solve()
for i, j in enumerate(match.pairA):
    if j == -1:
        continue
    if i > j:
        i, j = j, i
    if i + 1 == j:
        x, y = coo(i)
        grid[x][y] = ">"
        x, y = coo(j)
        grid[x][y] = "<"
    else:
        x, y = coo(i)
        grid[x][y] = "v"
        x, y = coo(j)
        grid[x][y] = "^"

print(ans)
for g in grid:
    print("".join(g))