Problem p02998 (python)

Fixed submission:
import sys
input = sys.stdin.readline
inpl = lambda: list(map(int,input().split()))
from collections import defaultdict

class UnionFind:
    def __init__(self, N=None):
        if N is None or N < 1:
            self.parent = defaultdict(lambda: -1)
        else:
            self.parent = [-1]*int(N)

    def root(self, n):
        if self.parent[n] < 0:
            return n
        else:
            m = self.root(self.parent[n])
            self.parent[n] = m
            return m

    def merge(self, m, n):
        rm = self.root(m)
        rn = self.root(n)
        if rm != rn:
            if -self.parent[rm] < -self.parent[rn]:
                rm, rn = rn, rm
            self.parent[rm] += self.parent[rn]
            self.parent[rn] = rm

    def size(self, n):
        return -self.parent[self.root(n)]
    
    def connected(self, m, n):
        return self.root(m) == self.root(n)
    
    def groups(self):
        if isinstance(self.parent,list):
            return list(map(lambda x: x<0, self.parent)).count(True)
        else: # self.parent: defaultdict
            return list(map(lambda x: x<0, self.parent.values())).count(True) 

class CountUp:
    def __init__(self, start=0):
        self.index = start-1

    def __call__(self):
        self.index += 1
        return self.index

Xi = defaultdict(CountUp())
Yi = defaultdict(CountUp())
N = int(input())
xy = []
for _ in range(N):
    x, y = inpl()
    xi = Xi[x]
    yi = Yi[y]
    xy.append((xi,yi))
Nx = len(Xi)
Ny = len(Yi)
Nxy = Nx + Ny
uf = UnionFind(Nx+Ny)
for xi, yi in xy:
    uf.merge(xi,yi+Nx)

gi = defaultdict(CountUp())
Ng = uf.groups()
edges = [0]*Ng
nx = [0]*Ng
ny = [0]*Ng
for i in range(Nx):
    nx[gi[uf.root(i)]] += 1
for i in range(Nx,Nxy):
    ny[gi[uf.root(i)]] += 1
for xi, yi in xy:
    edges[gi[uf.root(xi)]] += 1

ans = 0
for g in range(Ng):
    ans += nx[g]*ny[g] - edges[g]
print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline
inpl = lambda: list(map(int,input().split()))
from collections import defaultdict
class UnionFind:
    def __init__(self, N=None):
        if N is None or N < 1:
            self.parent = defaultdict(lambda: -1)
            self.children = defaultdict(list)
        else:
            self.parent = [-1]*int(N)
            self.children = [ [] for i in range(int(N)) ]

    def root(self, n):
        if self.parent[n] < 0:
            return n
        else:
            m = self.root(self.parent[n])
            self.parent[n] = m
            self.children[m].extend(self.children[n])
            self.children[n] = []
            return m

    def merge(self, m, n):
        rm = self.root(m)
        rn = self.root(n)
        if rm != rn:
            if -self.parent[rm] < -self.parent[rn]:
                rm, rn = rn, rm
            self.parent[rm] += self.parent[rn]
            self.parent[rn] = rm
            self.children[rm].append(rn)
            self.children[rm].extend(self.children[rn])
            self.children[rn] = []

    def size(self, n):
        return -self.parent[self.root(n)]
    
    def connected(self, m, n):
        return self.root(m) == self.root(n)
    
    def groups(self):
        if isinstance(self.parent,list):
            return list(map(lambda x: x<0, self.parent)).count(True)
        else: # self.parent: defaultdict
            return list(map(lambda x: x<0, self.parent.values())).count(True) 

    def collect_descendants(self, n):
        des = []
        for c in self.children[n]:
            self.collect_descendants(c)
            des.extend(self.children[c])
            self.children[c] = []
        self.children[n].extend(des)

    def all_groups(self):
        ans = []
        for r in range(N):
            if self.parent[r] >= 0:
                continue
            else:
                self.collect_descendants(r)
                ans.append( [r] + self.children[r] )
        return ans

class CountUp:
    def __init__(self, start=0):
        self.index = start-1

    def __call__(self):
        self.index += 1
        return self.index

Xi = defaultdict(CountUp())
Yi = defaultdict(CountUp())
N = int(input())
xy = []
for _ in range(N):
    x, y = inpl()
    xi = Xi[x]
    yi = Yi[y]
    xy.append((xi,yi))
Nx = len(Xi)
Ny = len(Yi)
#print(Nx,Ny)
uf = UnionFind(Nx+Ny)
for xi, yi in xy:
    uf.merge(xi,yi+Nx)
groups = uf.all_groups()
#print(uf.parent,uf.children)
#print(groups)
Ng = len(groups)
edge_nums = [0] * Ng
parents_id = dict()
for i in range(Ng):
    parents_id[groups[i][0]] = i
groups_xy = []
for g in groups:
    xs = []
    ys = []
    for k in g:
        if k < Nx:
            xs.append(k)
        else:
            ys.append(k-Nx)
    groups_xy.append([xs,ys])
for i in range(N):
    p = uf.root(xy[i][0])
    pi = parents_id[p]
    edge_nums[pi] += 1
ans = 0
for i in range(Ng):
    nx = len(list(filter(lambda s: s<Nx, groups[i])))
    ny = len(groups[i]) - nx
    ans += nx*ny - edge_nums[i]

print(ans)