Problem p02998 (python)

Fixed submission:
#!/usr/bin/env python3
import sys
from itertools import product
sys.setrecursionlimit(10**8)


INF = float("inf")


class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, 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 = 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):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, 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 compress(X, Y):
    dx = list(sorted(set(X)))
    dx_inv = {x: i for i, x in enumerate(dx)}
    dy = list(sorted(set(Y)))
    dy_inv = {y: i for i, y in enumerate(dy)}
    return (dx, dy), (dx_inv, dy_inv)


def main():
    N = int(input())
    X = [0]*N
    Y = [0]*N
    for i in range(N):
        X[i], Y[i] = map(int, input().split())

    (decx, decy), (encx, ency) = compress(X, Y)

    H = len(decx)
    uf = UnionFind(len(decx) + len(decy))

    for i in range(N):
        x, y = encx[X[i]], ency[Y[i]]+H
        # print(x, y)
        uf.union(x, y)

    dx = {}
    for i in range(len(decx)):
        key = uf.find(i)
        if key not in dx:
            dx[key] = 1
        else:
            dx[key] += 1
    dy = {}
    for i in range(H, H+len(decy)):
        key = uf.find(i)
        if key not in dy:
            dy[key] = 1
        else:
            dy[key] += 1
    # print(dx, dy)
    # print(uf.all_group_members())
    ans = 0
    for key in dx:
        if key in dy:
            ans += dy[key]*dx[key]
    print(ans-N)
    return


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
from itertools import product
sys.setrecursionlimit(10**8)


INF = float("inf")


def compress(X, Y):
    dx = list(sorted(set(X)))
    dx_inv = {x: i for i, x in enumerate(dx)}
    dy = list(sorted(set(Y)))
    dy_inv = {y: i for i, y in enumerate(dy)}
    return (dx, dy), (dx_inv, dy_inv)


def main():
    N = int(input())
    X = [0]*N
    Y = [0]*N
    for i in range(N):
        X[i], Y[i] = map(int, input().split())

    (dx, dy), (dx_inv, dy_inv) = compress(X, Y)

    X2Y = [set() for _ in dx]   # XからYを求める
    Y2X = [set() for _ in dy]   # YからXを求める
    for i in range(N):
        X2Y[dx_inv[X[i]]].add(dy_inv[Y[i]])
        Y2X[dy_inv[Y[i]]].add(dx_inv[X[i]])

    cnt = 0
    head = 0
    queue = [(dx_inv[X[i]], dy_inv[Y[i]]) for i in range(N)]
    while head < len(queue):    # 最終的な節点の個数だけのループ最悪(10^10/4)
        x, y = queue[head]
        head += 1
        samex = X2Y[x]
        samey = Y2X[y]
        # 全ての組み合わせを追加
        additive = []
        for nx in samex:
            if nx == x:
                continue
            nys = samey-X2Y[nx]
            additive.extend([(nx, ny) for ny in nys])
            cnt += len(nys)
            X2Y[nx] |= nys
            for ny in nys:
                Y2X[ny].add(nx)
    print(cnt)
    return


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
#!/usr/bin/env python3
import sys
from itertools import product


INF = float("inf")


def compress(X, Y):
    dx = list(sorted(set(X)))
    dx_inv = {x: i for i, x in enumerate(dx)}
    dy = list(sorted(set(Y)))
    dy_inv = {y: i for i, y in enumerate(dy)}
    return (dx, dy), (dx_inv, dy_inv)


# X = [-INF, 0, INF] + A + B + D
# Y = [-INF, 0, INF] + C + E + F
# (dx, dy), (dx_inv, dy_inv) = compress(X, Y)
# H = len(dy)
# W = len(dx)

def main():
    N = int(input())
    X = [0]*N
    Y = [0]*N
    for i in range(N):
        X[i], Y[i] = map(int, input().split())

    (dx, dy), (dx_inv, dy_inv) = compress(X, Y)

    X2Y = [set() for _ in dx]   # XからYを求める
    Y2X = [set() for _ in dy]   # YからXを求める
    for i in range(N):
        X2Y[dx_inv[X[i]]].add(dy_inv[Y[i]])
        Y2X[dy_inv[Y[i]]].add(dx_inv[X[i]])

    cnt = 0
    head = 0
    queue = [(dx_inv[X[i]], dy_inv[Y[i]]) for i in range(N)]
    while head < len(queue):    # 最終的な節点の個数だけのループ最悪(10^10/4)
        x, y = queue[head]
        head += 1
        samex = X2Y[x]
        samey = Y2X[y]
        # 全ての組み合わせを追加
        additive = []
        for nx in samex:
            if nx == x:
                continue
            nys = samey-X2Y[nx]
            additive.extend([(nx, ny) for ny in nys])
            cnt += len(nys)
            X2Y[nx] |= nys
            for ny in nys:
                Y2X[ny].add(nx)
    print(cnt)
    return


if __name__ == '__main__':
    main()