Problem p03108 (python)

Fixed submission:

from heapq import heappush, heappop, heapify
from collections import deque, defaultdict, Counter
import itertools
from itertools import permutations, combinations
import sys
import bisect
import string
import math
import time
from fractions import gcd
#import random


def I(): return int(input())


def MI(): return map(int, input().split())


def LI(): return [int(i) for i in input().split()]


def LI_(): return [int(i)-1 for i in input().split()]


def StoI(): return [ord(i)-97 for i in input()]


def ItoS(nn): return chr(nn+97)


def input(): return sys.stdin.readline().rstrip()


def show(*inp, end='\n'):
    if show_flg:
        print(*inp, end=end)


YN = ['Yes', 'No']
MOD = 10**9+7
inf = float('inf')
l_alp = string.ascii_lowercase
u_alp = string.ascii_uppercase
ts = time.time()
sys.setrecursionlimit(10**6)


# show_flg = True
show_flg = False


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

    def find(self, x):
        # ret = x
        # stack = [x]
        # while True:
        #     data = self.parents[stack[-1]]
        #     if data < 0:
        #         break
        #     stack.append(data)

        # while len(stack) > 1:
        #     ret = stack.pop()
        #     self.parents[data] = data

        # return ret

        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 main():
    N, M = MI()
    # node_list = [UnionFindNode(i) for i in range(N+1)]
    uf = UnionFind(N+1)
    U = [0] * M
    V = [0] * M
    ans = [0] * M
    ans[M-1] = N * (N - 1) // 2

    for i in range(M):
        U[i], V[i] = MI()

    for i in range(M-1, 0, -1):
        u = U[i]
        v = V[i]

        if uf.same(u, v):
            ans[i-1] = ans[i]
        else:
            ans[i-1] = ans[i] - uf.size(u) * uf.size(v)
            uf.union(u, v)

    print(*ans)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):

from heapq import heappush, heappop, heapify
from collections import deque, defaultdict, Counter
import itertools
from itertools import permutations, combinations
import sys
import bisect
import string
import math
import time
from fractions import gcd
#import random


def I(): return int(input())


def MI(): return map(int, input().split())


def LI(): return [int(i) for i in input().split()]


def LI_(): return [int(i)-1 for i in input().split()]


def StoI(): return [ord(i)-97 for i in input()]


def ItoS(nn): return chr(nn+97)


def input(): return sys.stdin.readline().rstrip()


def show(*inp, end='\n'):
    if show_flg:
        print(*inp, end=end)


YN = ['Yes', 'No']
MOD = 10**9+7
inf = float('inf')
l_alp = string.ascii_lowercase
u_alp = string.ascii_uppercase
ts = time.time()
sys.setrecursionlimit(10**6)


# show_flg = True
show_flg = False


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

    def find(self, x):
        ret = x
        stack = [x]
        while True:
            data = self.parents[stack[-1]]
            if data < 0:
                break
            stack.append(data)

        while len(stack) > 1:
            ret = stack.pop()
            self.parents[data] = data

        return ret

        # if self.parents[x] < 0:
        #     print(x, self.parents[x])
        #     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 main():
    N, M = MI()
    # node_list = [UnionFindNode(i) for i in range(N+1)]
    uf = UnionFind(N+1)
    U = [0] * M
    V = [0] * M
    ans = [0] * M
    ans[M-1] = N * (N - 1) // 2

    for i in range(M):
        U[i], V[i] = MI()

    for i in range(M-1, 0, -1):
        u = U[i]
        v = V[i]

        if uf.same(u, v):
            ans[i-1] = ans[i]
        else:
            ans[i-1] = ans[i] - uf.size(u) * uf.size(v)
            uf.union(u, v)

    print(*ans)


if __name__ == '__main__':
    main()