Problem p02998 (python)

Fixed submission:
import sys
from collections import defaultdict
def input(): return sys.stdin.readline().strip()

class UnionFind():
    """
    https://note.nkmk.me/python-union-find/
    DFSの上位互換と考えて良い
    ２要素x, yがpath-connectedかどうかをlogオーダーで判定する（螺旋本の14.1節参照）
    さらに連結成分の要素数がO(1)で取得可能なように改造してある
    """
    def __init__(self, n):
        """
        要素数をnとして、各ノードを0,1,...,(n-1)の番号で管理する
        parentsは各ノードの属する木の根を表す
        ただし根ノードのparentには(その木のノード数)*(-1)を格納する
        """
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        """
        xの属する木の根を返す
        このとき同時に経路圧縮して、探索途中のノードを全て根に繋ぎ直す
        """
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        """
        x, yのそれぞれ属する木Tx, Tyの根同士を繋ぐ
        このとき木の要素数が小さい方を大きい方に繋ぐ（rankではなくsizeを用いる）
        """
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        """
        xの属する木の要素数を返す
        根の親を要素数の(-1)倍で定めておいたおかげでO(1)で取得可能
        """
        return -self.parents[self.find(x)]

    def same(self, x, y):
        """
        xとyがpath-connectedかを判定する
        """
        return self.find(x) == self.find(y)

    def members(self, x):
        """
        xの属する木の要素を列挙する
        """
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        """
        連結成分の代表元のリストを返す
        """
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        """
        連結成分の個数を返す
        """
        return len(self.roots())

    def all_group_members(self):
         """
         連結成分およびそれぞれの代表元をまとめた辞書を返す
         代表元がキーになってる
         """
         return {r: self.members(r) for r in self.roots()}

    def __str__(self):
        """
        連結成分およびその代表元を出力
        """
        return '\n'.join('{}: {}'.format(r, self.members(r)) for r in self.roots())



def main():

    """
    x方向とy方向に分けて頂点を辺だと思うのは定石？？？（確かに長方形がメインなら直積分解したくなる？）
    このとき問題文は「長さ３のパスがあればパスを１本追加する操作を繰り返して、各連結成分を完全二部グラフにせよ」

    完全二部グラフになるのは背理法からすぐわかる。なので解法は(¥Sum_{i:component}X_i Y_i) - N
    """

    N = int(input())
    vertex_x = set([])
    vertex_y = set([])
    edge = []
    for _ in range(N):
        x, y = map(int, input().split())
        vertex_x.add(x)
        vertex_y.add(y)
        edge.append((x, y))

    # UnionFind()の都合上頂点は0-indexedの番号列で管理したいので
    # 頂点に対して連番づけを行う
    d = defaultdict(int)
    num = 0
    for x in vertex_x:
        d[(0, x)] = num
        num += 1
    x_num = num
    for y in vertex_y:
        d[(1, y)] = num
        num += 1

    # ここからメイン
    uftree = UnionFind(num)
    for x, y in edge:
        uftree.union(d[(0, x)], d[(1, y)])

    ans = 0
    for r in uftree.roots():
        X = 0
        Y = 0
        for v in uftree.members(r):
            if v < x_num:
                X += 1
            else:
                Y += 1
        ans += X * Y
    print(ans - N)



if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
import sys
def input(): return sys.stdin.readline().strip()

def main():

    """
    <UnionFindを使わない別解>

    x方向とy方向に分けて頂点を辺だと思うのは定石？？？（確かに長方形がメインなら直積分解したくなる？）
    このとき問題文は「長さ３のパスがあればパスを１本追加する操作を繰り返して、各連結成分を完全二部グラフにせよ」

    完全二部グラフになるのは背理法からすぐわかる。なので解法は(¥Sum_{i:component}X_i Y_i) - N
    """

    N = int(input())
    V = 10 ** 5 + 1 # padding for y
    repn = [[] for _ in range(V * 2)]
    for _ in range(N):
        x, y = map(int, input().split())
        repn[x].append(V + y)
        repn[V + y].append(x)

    visited = [0] * (V * 2)
    count = [0, 0]

    # dfsでその都度連結成分を求める
    def dfs(v):
        if visited[v]:
            return
        visited[v] = 1
        count[v // V] += 1
        for u in repn[v]:
            dfs(u)

    ans = 0
    for i in range(V * 2):
        if visited[i]:
            continue
        count[:] = [0, 0]
        dfs(i)
        ans += count[0] * count[1]
    print(ans - N)


if __name__ == "__main__":
    main()