Problem p03108 (Python) — buggy submission, verdict: Time Limit Exceeded

class UnionFind(object):
    def __init__(self, n=1):
        self.par = [i for i in range(n)]  # グループ番号
        self.rank = [0 for _ in range(n)]
        self.num = [1 for _ in range(n)]

    def find(self, x):
        """
        x が属するグループを探索
        """
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def union(self, x, y):
        """
        x と y のグループを結合
        """
        x = self.find(x)
        y = self.find(y)
        if x != y:
            if self.rank[x] < self.rank[y]:
                x, y = y, x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1
            self.par[y] = x
            self.num[x] += self.num[y]
            self.num[y] = 0

    def is_same(self, x, y):
        """
        x と y が同じグループか否か
        """
        return self.find(x) == self.find(y)


N, M = map(int, input().split())
L = [list(map(int, input().split())) for i in range(M)]
A = [0] * M
D = N * (N - 1) // 2
u = UnionFind(N)
a = 0
for i in reversed(range(M)):
    # u.union(L[i][0]-1, L[i][1]-1)
    # a = 0
    # for j in range(N):
    #     # print(i, j, u.num[j])
    #     if u.num[j] >= 2:
    #         a += u.num[j] * (u.num[j] - 1) // 2
    A[i] = a
    # print(A[i])
    if u.find(L[i][0] - 1) != u.find(L[i][1] - 1):
        if u.num[u.find(L[i][0] - 1)] > 1:
            a -= u.num[u.find(L[i][0] - 1)] * (u.num[u.find(L[i][0] - 1)] - 1) // 2
        if u.num[u.find(L[i][1] - 1)] > 1:
            a -= u.num[u.find(L[i][1] - 1)] * (u.num[u.find(L[i][1] - 1)] - 1) // 2
        u.union(L[i][0] - 1, L[i][1] - 1)
        # if u.find(L[i][0] - 1) != u.find(L[i][1] - 1):
        a += u.num[u.find(L[i][0] - 1)] * (u.num[u.find(L[i][0] - 1)] - 1) // 2

for i in range(M):
    print(D - A[i])


Error: TLE

Fixed submission:

class UnionFind(object):
    def __init__(self, n=1):
        self.par = [i for i in range(n)]  # グループ番号
        self.rank = [0 for _ in range(n)]
        self.num = [1 for _ in range(n)]

    def find(self, x):
        """
        x が属するグループを探索
        """
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def union(self, x, y):
        """
        x と y のグループを結合
        """
        x = self.find(x)
        y = self.find(y)
        if x != y:
            if self.rank[x] < self.rank[y]:
                x, y = y, x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1
            self.par[y] = x
            self.num[x] += self.num[y]
            self.num[y] = 0

    def is_same(self, x, y):
        """
        x と y が同じグループか否か
        """
        return self.find(x) == self.find(y)


N, M = map(int, input().split())
L = [list(map(int, input().split())) for i in range(M)]
A = [0] * M
D = N * (N - 1) // 2
u = UnionFind(N)
for i in reversed(range(M)):
    # u.union(L[i][0]-1, L[i][1]-1)
    a = 0
    for j in range(N):
        # print(i, j, u.num[j])
        if u.num[j] >= 2:
            a += u.num[j] * (u.num[j] - 1) // 2
    A[i] = a
    # print(A[i])
    u.union(L[i][0] - 1, L[i][1] - 1)
for i in range(M):
    print(D - A[i])