Problem p02558 (python)

Fixed submission:
class UnionFind():
    "Either 0-indexed or 1-indexed"

    __slots__ = ["parent"]

    def __init__(self, size):
        self.parent = [-1] * (size + 1)
    
    def find(self, a):
        path = []
        while self.parent[a] > 0:
            path.append(a)
            a = self.parent[a]
        for child in path:
            self.parent[child] = a
        return a

    def union(self, a, b):
        a = self.find(a)
        b = self.find(b)

        if a == b:
            return
        else:
            if self.parent[a] == self.parent[b]:
                self.parent[a] = b
                self.parent[b] -= 1
            elif self.parent[a] < self.parent[b]: #aのほうが大きい
                self.parent[b] = a
            else:
                self.parent[a] = b #bのほうが大きい
        
    def same(self, a, b):
        return self.find(a) == self.find(b)


def main():
    import sys
    input = sys.stdin.buffer.readline
    N, Q = map(int, input().split())
    tuvs = map(int, sys.stdin.buffer.read().split())
    ans = []
    uf = UnionFind(N)
    for t, u, v in zip(tuvs, tuvs, tuvs):
        if t:
            ans.append('1' if uf.same(u, v) else '0')
        else:
            uf.union(u, v)

    print('\n'.join(ans))


if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
# https://atcoder.jp/contests/practice2/submissions/16598132
# 上記のリンクの提出をいじらせてもらってます。

class mf_graph:

    __slots__ = ["N", "graph", "capacities", "level", "it"]

    def __init__(self, N):
        self.N = N
        self.graph = [[] for _ in range(N)]
        self.capacities = [dict() for _ in range(N)]

    def add_edge(self, v, w, cap):
        self.graph[v].append(w)
        self.graph[w].append(v)
        self.capacities[v][w] = cap
        self.capacities[w][v] = 0

    # def bfs(self, s, t):
    #     self.level = [-1] * self.N
    #     q = deque([s])
    #     self.level[s] = 0
    #     while q:
    #         v = q.popleft()
    #         if v == t:
    #             return True
    #         for w, cap in self.capacities[v].items():
    #             if cap and self.level[w] == -1:
    #                 self.level[w] = self.level[v] + 1
    #                 q.append(w)
    #     return False

    def bfs(self, s, t):
        self.level = [-1] * self.N
        q = [s]
        nq = []
        self.level[s] = 0
        while q:
            nq = []
            for v in q:
                if v == t:
                    return True
                for w, cap in self.capacities[v].items():
                    if cap and self.level[w] == -1:
                        self.level[w] = self.level[v] + 1
                        nq.append(w)
            q = nq
        return False

    def dfs(self, s, t, up):
        st = [t]
        while st:
            v = st[-1]
            if v == s:
                flow = up
                for i in range(len(st) - 2):
                    flow = min(flow, self.capacities[st[i+1]][st[i]])
                for i in range(len(st) - 1):
                    self.capacities[st[i]][st[i+1]] += flow
                    self.capacities[st[i+1]][st[i]] -= flow
                return flow
            while self.it[v] < len(self.graph[v]):
                w = self.graph[v][self.it[v]]
                cap = self.capacities[w][v]
                if cap and self.level[w] != -1:
                    if self.level[v] > self.level[w]:
                        st.append(w)
                        break
                self.it[v] += 1
            else:
                st.pop()
                self.level[v] = self.N
        return 0

    def flow(self, s, t, flow_limit=18446744073709551615):
        flow = 0
        while flow < flow_limit and self.bfs(s, t):
            self.it = [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
        q = [s]
        while q:
            v = q.pop()
            visited[v] = True
            for w, cap in self.capacities[v].items():
                if cap and not visited[w]:
                    q.append(w)


# """
N, M = map(int, input().split())
S = [list(input()) for i in range(N)]
"""
with open("flow.txt") as f:
    N, M = map(int, f.readline().rstrip().split())
    S = [list(f.readline().rstrip()) for i in range(N)]
#"""
dc = mf_graph(N*M+2)

s = N*M
t = N*M+1
for i in range(N):
    for j in range(M):
        if (i+j) % 2 == 0:
            dc.add_edge(s, M*i+j, 1)
        else:
            dc.add_edge(M*i+j, t, 1)

for i in range(N):
    for j in range(M):
        if j+1 < M and S[i][j] == "." and S[i][j+1] == ".":
            u, v = M*i+j, M*i+j+1
            if (i+j) % 2 == 1:
                u, v = v, u
            dc.add_edge(u, v, 1)
        if i+1 < N and S[i][j] == "." and S[i + 1][j] == ".":
            u, v = M*i+j, M*(i+1)+j
            if (i+j) % 2 == 1:
                u, v = v, u
            dc.add_edge(u, v, 1)

print(dc.flow(s, t))
for u in range(N*M+2):
    ui, uj = divmod(u, M)
    if (ui+uj) % 2 == 1:
        continue
    for v, cap in dc.capacities[u].items():
        if cap != 0:
            continue
        vi, vj = divmod(v, M)
        if u != s and u != t and v != s and v != t:
            if ui+1 == vi:
                S[ui][uj] = "v"
                S[vi][vj] = "^"
            elif ui == vi+1:
                S[ui][uj] = "^"
                S[vi][vj] = "v"
            elif uj+1 == vj:
                S[ui][uj] = ">"
                S[vi][vj] = "<"
            elif uj == vj+1:
                S[ui][uj] = "<"
                S[vi][vj] = ">"

for res in S:
    print("".join(res))

Buggy submission (verdict: Runtime Error):
# import cProfile

# https://atcoder.jp/contests/practice2/submissions/16598132
# 上記のリンクの提出をいじらせてもらってます。

# from collections import deque


class mf_graph:

    # __slots__ = ["N", "graph", "capacities", "level", "it"]

    def __init__(self, N):
        self.N = N
        self.graph = [[] for _ in range(N)]
        self.capacities = [dict() for _ in range(N)]

    def add_edge(self, v, w, cap):
        self.graph[v].append(w)
        self.graph[w].append(v)
        self.capacities[v][w] = cap
        self.capacities[w][v] = 0

    # def bfs(self, s, t):
    #     self.level = [-1] * self.N
    #     q = deque([s])
    #     self.level[s] = 0
    #     while q:
    #         v = q.popleft()
    #         if v == t:
    #             return True
    #         for w, cap in self.capacities[v].items():
    #             if cap and self.level[w] == -1:
    #                 self.level[w] = self.level[v] + 1
    #                 q.append(w)
    #     return False

    def bfs(self, s, t):
        self.level = [-1] * self.N
        q = [s]
        nq = []
        self.level[s] = 0
        while q:
            nq.clear()
            for v in q:
                if v == t:
                    return True
                for w, cap in self.capacities[v].items():
                    if cap and self.level[w] == -1:
                        self.level[w] = self.level[v] + 1
                        nq.append(w)
            q = nq
        return False

    def dfs(self, s, t, up):
        st = [t]
        while st:
            v = st[-1]
            if v == s:
                flow = up
                for i in range(len(st) - 2):
                    flow = min(flow, self.capacities[st[i+1]][st[i]])
                for i in range(len(st) - 1):
                    self.capacities[st[i]][st[i+1]] += flow
                    self.capacities[st[i+1]][st[i]] -= flow
                return flow
            while self.it[v] < len(self.graph[v]):
                w = self.graph[v][self.it[v]]
                cap = self.capacities[w][v]
                if cap and self.level[w] != -1:
                    if self.level[v] > self.level[w]:
                        st.append(w)
                        break
                self.it[v] += 1
            else:
                st.pop()
                self.level[v] = self.N
        return 0

    def flow(self, s, t, flow_limit=float('inf')):
        flow = 0
        while flow < flow_limit and self.bfs(s, t):
            self.it = [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
        q = deque([s])
        while q:
            v = q.pop()
            visited[v] = True
            for w, cap in self.capacities[v].items():
                if cap and not visited[w]:
                    q.append(w)

#"""
N, M = map(int, input().split())
S = [list(input()) for i in range(N)]
"""
with open("flow.txt") as f:
    N, M = map(int, f.readline().rstrip().split())
    S = [list(f.readline().rstrip()) for i in range(N)]
#"""
dc = mf_graph(N*M+2)

s = N*M
t = N*M+1
for i in range(N):
    for j in range(M):
        if (i+j) % 2 == 0:
            dc.add_edge(s, M*i+j, 1)
        else:
            dc.add_edge(M*i+j, t, 1)

for i in range(N):
    for j in range(M):
        if j+1 < M and S[i][j] == "." and S[i][j+1] == ".":
            u, v = M*i+j, M*i+j+1
            if (i+j) % 2 == 1:
                u, v = v, u
            dc.add_edge(u, v, 1)
        if i+1 < N and S[i][j] == "." and S[i + 1][j] == ".":
            u, v = M*i+j, M*(i+1)+j
            if (i+j) % 2 == 1:
                u, v = v, u
            dc.add_edge(u, v, 1)

print(dc.flow(s, t))
for u in range(N*M+2):
    ui, uj = divmod(u, M)
    if (ui+uj) % 2 == 1:
        continue
    for v, cap in dc.capacities[u].items():
        if cap != 0:
            continue
        vi, vj = divmod(v, M)
        if u != s and u != t and v != s and v != t:
            if ui+1 == vi:
                S[ui][uj] = "v"
                S[vi][vj] = "^"
            elif ui == vi+1:
                S[ui][uj] = "^"
                S[vi][vj] = "v"
            elif uj+1 == vj:
                S[ui][uj] = ">"
                S[vi][vj] = "<"
            elif uj == vj+1:
                S[ui][uj] = "<"
                S[vi][vj] = ">"

for res in S:
    print("".join(res))