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

import sys
from itertools import combinations

input = sys.stdin.readline


class UnionFind:
    """Union Find class.

    "Path compression" and "Union by rank" are used.

    References:
        <https://en.wikipedia.org/wiki/Disjoint-set_data_structure>
    """

    def __init__(self, N):
        self.N = N
        self.__make_set()

    def __make_set(self):
        self._parent = list(range(self.N + 1))
        self._rank = [0] * (self.N + 1)
        self._size = [1] * (self.N + 1)

    def find(self, x):
        if self._parent[x] != x:
            self._parent[x] = self.find(self._parent[x])
        return self._parent[x]

    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)

        if x_root == y_root:
            return

        x_rank = self._rank[x_root]
        y_rank = self._rank[y_root]
        if x_rank > y_rank:
            self._parent[y_root] = x_root
            self._size[x_root] += self._size[y_root]
        elif x_rank < y_rank:
            self._parent[x_root] = y_root
            self._size[y_root] += self._size[x_root]
        else:
            self._parent[y_root] = x_root
            self._rank[x_root] += 1
            self._size[x_root] += self._size[y_root]

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

    def size(self, x):
        return self._size[self.find(x)]


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

    uf = UnionFind(N)
    ans = [0] * M
    ans[-1] = N * (N - 1) // 2
    for i in reversed(range(M - 1)):
        a, b = A[i + 1], B[i + 1]
        if uf.same_set(a, b):
            ans[i] = ans[i + 1]
        else:
            ans[i] = ans[i + 1] - uf.size(a) * uf.size(b)
        uf.union(a, b)

    print("\\n".join(map(str, ans)))


if __name__ == "__main__":
    main()


Error: TLE

Fixed submission:

import sys
from itertools import combinations

input = sys.stdin.readline


class UnionFind:
    """Union Find class.

    "Path compression" and "Union by rank" are used.

    References:
        <https://en.wikipedia.org/wiki/Disjoint-set_data_structure>
    """

    def __init__(self, N):
        self.N = N
        self.__make_set()
        self.elem = {i: set([i]) for i in range(1, N + 1)}

    def __make_set(self):
        self.parent = list(range(self.N + 1))
        self.rank = [0] * (self.N + 1)

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        x_root = self.find(x)
        y_root = self.find(y)

        if x_root == y_root:
            return

        x_rank = self.rank[x_root]
        y_rank = self.rank[y_root]
        if x_rank > y_rank:
            self.parent[y_root] = x_root
            self.union_elem(y_root, x_root)
        elif x_rank < y_rank:
            self.parent[x_root] = y_root
            self.union_elem(x_root, y_root)
        else:
            self.parent[y_root] = x_root
            self.rank[x_root] += 1
            self.union_elem(y_root, x_root)

    def union_elem(self, src_root, dst_root):
        self.elem[dst_root] |= self.elem[src_root]
        del self.elem[src_root]

    def elem_counts(self):
        return [len(s) for s in self.elem.values()]


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

    uf = UnionFind(N)
    ans = [0] * M
    for i in reversed(range(M)):
        ans[i] = sum(a * b for a, b in combinations(uf.elem_counts(), r=2))
        uf.union(A[i], B[i])

    print("\\n".join(map(str, ans)))


if __name__ == "__main__":
    main()