Problem p03157 (python)

Fixed submission:
class UnionFind:
    def __init__(self, n):
        self.v = [-1 for _ in range(n)]  # 根(負): 連結頂点数 * (-1) / 子(正): 根の頂点番号(0-indexed)
        self.bk = [1 if s[r][c] == '#' else 0 for r in range(H) for c in range(W)]
        self.wt = [1 if s[r][c] == '.' else 0 for r in range(H) for c in range(W)]

    def find(self, x):  # xを含む木における根の頂点番号を返す
        if self.v[x] < 0:  # (負)は根
            return x
        else:  # 根の頂点番号
            self.v[x] = self.find(self.v[x])  # uniteでは, 旧根に属する頂点の根が旧根のままなので更新
            return self.v[x]

    def unite(self, x, y):  # 違う根に属していたらrankが低くなるように連結
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return
        if -self.v[x] < -self.v[y]:  # size比較, 　(-1) * (連結頂点数 * (-1)), (正)同士の大小比較
            x, y = y, x  # 連結頂点数が少ない方をyにすると, findでの更新回数が減る？
        self.v[x] += self.v[y]  # 連結頂点数の和を取る, 連結頂点数 * (-1)
        self.v[y] = x  # 連結頂点数が少ないy(引数yの根の頂点番号)の根をx(引数xの根の頂点番号)にする

        self.bk[x % (H * W)] += self.bk[y % (H * W)]
        self.wt[x % (H * W)] += self.wt[y % (H * W)]

    def root(self, x):
        return self.v[x] < 0  # (負)は根

    def same(self, x, y):
        return self.find(x) == self.find(y)  # 同じ根に属するか

    def size(self, x):
        return -self.v[self.find(x)]  # 連結頂点数を返す

    def calc(self, x):
        return self.bk[x % (H * W)] * self.wt[x % (H * W)]

dr = [0, 0, 1, -1]
dc = [1, -1, 0, 0]

H, W = map(int, input().split())
s = [input() for _ in range(H)]
uf = UnionFind(2 * H * W)  # 全頂点 * 2
for r in range(H):
    for c in range(W):
        if s[r][c] == '#':
            for i in range(4):
                nr = r + dr[i]
                nc = c + dc[i]
                if 0 <= nr < H and 0 <= nc < W:
                    if s[nr][nc] == '.':
                        uf.unite(r * W + c, H * W + (nr * W + nc))  # #: *1, .: * 2
        else:
            for i in range(4):
                nr = r + dr[i]
                nc = c + dc[i]
                if 0 <= nr < H and 0 <= nc < W:
                    if s[nr][nc] == '#':
                        uf.unite(H * W + (r * W + c), nr * W + nc)  # #: *1, .: * 2

ans = 0
for r in range(H):
    for c in range(W):
        if uf.root(r * W + c):
            # print(r, c, '#')
            sz = uf.size(r * W + c)
            # print(' ', r * W + c, sz, '#')
            if sz > 1:
                ans += uf.calc(r * W + c)

        if uf.root(H * W + (r * W + c)):
            # print(r, c, '.')
            sz = uf.size(H * W + (r * W + c))
            # print(' ', r * W + c, sz, '.')
            if sz > 1:
                ans += uf.calc(H * W + (r * W + c))

print(ans)
# print(uf.v)

Buggy submission (verdict: Runtime Error):
class UnionFind:
    def __init__(self, n):
        self.v = [-1 for _ in range(n)]  # 根(負): 連結頂点数 * (-1) / 子(正): 根の頂点番号(0-indexed)

    def find(self, x):  # xを含む木における根の頂点番号を返す
        if self.v[x] < 0:  # (負)は根
            return x
        else:  # 根の頂点番号
            self.v[x] = self.find(self.v[x])  # uniteでは, 旧根に属する頂点の根が旧根のままなので更新
            return self.v[x]

    def unite(self, x, y):  # 違う根に属していたらrankが低くなるように連結
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return
        if -self.v[x] < -self.v[y]:  # size比較, 　(-1) * (連結頂点数 * (-1)), (正)同士の大小比較
            x, y = y, x  # 連結頂点数が少ない方をyにすると, findでの更新回数が減る？
        self.v[x] += self.v[y]  # 連結頂点数の和を取る, 連結頂点数 * (-1)
        self.v[y] = x  # 連結頂点数が少ないy(引数yの根の頂点番号)の根をx(引数xの根の頂点番号)にする

    def root(self, x):
        return self.v[x] < 0  # (負)は根

    def same(self, x, y):
        return self.find(x) == self.find(y)  # 同じ根に属するか

    def size(self, x):
        return -self.v[self.find(x)]  # 連結頂点数を返す


dr = [0, 0, 1, -1]
dc = [1, -1, 0, 0]

H, W = map(int, input().split())
s = [input() for _ in range(H)]
uf = UnionFind(2 * H * W)  # 全頂点 * 2
for r in range(H):
    for c in range(W):
        if s[r][c] == '#':
            for i in range(4):
                nr = r + dr[i]
                nc = c + dc[i]
                if 0 <= nr < H and 0 <= nc < W:
                    if s[nr][nc] == '.':
                        uf.unite(r * W + c, (nr * W + nc) * 2)  # #: *1, .: * 2
        else:
            for i in range(4):
                nr = r + dr[i]
                nc = c + dc[i]
                if 0 <= nr < H and 0 <= nc < W:
                    if s[nr][nc] == '#':
                        uf.unite((r * W + c) * 2, nr * W + nc)  # #: *1, .: * 2

ans = 0
for r in range(H):
    for c in range(W):
        if uf.root(r * W + c):
            sz = uf.size(r * W + c)

            wt = sz // 2
            bk = sz - wt
            ans += wt * bk

        elif uf.root((r * W + c) * 2):
            sz = uf.size((r * W + c) * 2)

            bk = sz // 2
            wt = sz - wt
            ans += wt * bk

print(ans)