Problem p03440 (python)

Fixed submission:
# python template for atcoder1
import heapq
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline
# -*- coding: utf-8 -*-
# python template for atcoder1
sys.setrecursionlimit(10**9)
input = sys.stdin.readline


class UnionFind:
    """
    sizeによる実装
    """

    def __init__(self, N, L):
        self.parent = [i for i in range(N)]
        self.size = [1 for _ in range(N)]
        self.costs = [[l] for l in L]

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

    def union(self, x, y) -> "min cost":
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return
        if self.size[px] < self.size[py]:
            self.parent[px] = py
            self.size[py] += self.size[px]
            for new_mem in self.costs[px]:
                heapq.heappush(self.costs[py], new_mem)
            self.costs[px] = []
        else:
            self.parent[py] = px
            self.size[px] += self.size[py]
            for new_mem in self.costs[py]:
                heapq.heappush(self.costs[px], new_mem)
            self.costs[py] = []

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

    def get_parents(self) -> "list":
        return [i for i, p in enumerate(self.parent) if i == p]


N, M = map(int, input().split())
L = list(map(int, input().split()))
edges = [list(map(int, input().split())) for _ in range(M)]

Un = UnionFind(N, L)
for s, g in edges:
    Un.union(s, g)


if M == N-1:
    print(0)
elif N < 2*(N-M-1):
    print("Impossible")
else:
    cost = 0
    cnt = 0
    for l in Un.get_parents():
        cost += heapq.heappop(Un.costs[l])
        cnt += len(Un.costs[l])

    remain_l = []
    for l in Un.get_parents():
        for v in Un.costs[l]:
            heapq.heappush(remain_l, v)
    for _ in range(2*(N-M-1)-len(Un.get_parents())):
        cost += heapq.heappop(remain_l)
    print(cost)

Buggy submission (verdict: Runtime Error):
# python template for atcoder1
import heapq
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline
# -*- coding: utf-8 -*-
# python template for atcoder1
sys.setrecursionlimit(10**9)
input = sys.stdin.readline


class UnionFind:
    """
    sizeによる実装
    """

    def __init__(self, N, L):
        self.parent = [i for i in range(N)]
        self.size = [1 for _ in range(N)]
        self.costs = [[l] for l in L]

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

    def union(self, x, y) -> "min cost":
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return
        if self.size[px] < self.size[py]:
            self.parent[px] = py
            self.size[py] += self.size[px]
            for new_mem in self.costs[px]:
                heapq.heappush(self.costs[py], new_mem)
            self.costs[px] = []
        else:
            self.parent[py] = px
            self.size[px] += self.size[py]
            for new_mem in self.costs[py]:
                heapq.heappush(self.costs[px], new_mem)
            self.costs[py] = []

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

    def merge(self, x, y):
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return 0
        if len(self.costs[px]) == 0 or len(self.costs[py]) == 0:
            print("Impossible")
            exit()
        cost = heapq.heappop(self.costs[px])+heapq.heappop(self.costs[py])
        self.union(px, py)
        return cost

    def get_parents(self) -> "list":
        return [i for i, p in enumerate(self.parent) if i == p]


N, M = map(int, input().split())
L = list(map(int, input().split()))
edges = [list(map(int, input().split())) for _ in range(M)]

Un = UnionFind(N, L)
for s, g in edges:
    Un.union(s, g)


if M == N-1:
    print(0)
elif N < 2*(N-M-1):
    print("Impossible")
else:
    cost = 0
    for l in Un.get_parents():
        cost += heapq.heappop(Un.costs[l])

    remain_l = []
    for l in Un.get_parents():
        for v in Un.costs[l]:
            heapq.heappush(remain_l, v)
    for _ in range(2*(N-M-1)):
        cost += heapq.heappop(remain_l)
    print(cost)

Buggy submission (verdict: Runtime Error):
# python template for atcoder1
import heapq
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline
# -*- coding: utf-8 -*-
# python template for atcoder1
sys.setrecursionlimit(10**9)
input = sys.stdin.readline


class UnionFind:
    """
    sizeによる実装
    """

    def __init__(self, N, L):
        self.parent = [i for i in range(N)]
        self.size = [1 for _ in range(N)]
        self.costs = [[l] for l in L]

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

    def union(self, x, y) -> "min cost":
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return
        if self.size[px] < self.size[py]:
            self.parent[px] = py
            self.size[py] += self.size[px]
            for new_mem in self.costs[px]:
                heapq.heappush(self.costs[py], new_mem)
            self.costs[px] = []
        else:
            self.parent[py] = px
            self.size[px] += self.size[py]
            for new_mem in self.costs[py]:
                heapq.heappush(self.costs[px], new_mem)
            self.costs[py] = []

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

    def merge(self, x, y):
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return 0
        if len(self.costs[px]) == 0 or len(self.costs[py]) == 0:
            print("Impossible")
            exit()
        cost = heapq.heappop(self.costs[px])+heapq.heappop(self.costs[py])
        self.union(px, py)
        return cost

    def get_parents(self) -> "list":
        return [i for i, p in enumerate(self.parent) if i == p]


N, M = map(int, input().split())
L = list(map(int, input().split()))
edges = [list(map(int, input().split())) for _ in range(M)]

Un = UnionFind(N, L)
for s, g in edges:
    Un.union(s, g)


cost = 0
for l in Un.get_parents():
    cost += heapq.heappop(Un.costs[l])

remain_l = []
for l in Un.get_parents():
    for v in Un.costs[l]:
        heapq.heappush(remain_l, v)
for _ in range(2*(N-M-1)):
    cost += heapq.heappop(remain_l)
print(cost)