Problem p03440 (python)

Fixed submission:
from collections import Counter, defaultdict
import sys
sys.setrecursionlimit(10 ** 5 + 10)
input = sys.stdin.readline
from math import factorial
import heapq, bisect
import math
import itertools


import queue
from collections import deque


parent_data = []

def union(ind):
    global parent_data
    if parent_data[ind] != ind:
        parent_data[ind] = union(parent_data[ind])
    return parent_data[ind]


def main():
    node_num, edge_num = map(int, input().split())
    cost_data = list(map(int, input().split()))
    data = [list(map(int, input().split())) for i in range(edge_num)]

    global parent_data
    parent_data = [i for i in range(node_num)]

    for i in range(edge_num):
        a, b = data[i]
        parent_a = union(a)
        parent_b = union(b)
        parent_data[min(parent_a, parent_b)] = max(parent_a, parent_b)

    parent_count = defaultdict(int)
    cost_heap = defaultdict(list)
    for i in range(node_num):
        parent_i = union(i)
        parent_count[parent_i] += 1
        heapq.heappush(cost_heap[parent_i], (cost_data[i], parent_i))

    parent_count = list(sorted(parent_count.items(), key=lambda x:-x[1]))

    ans = 0
    ori_ind, ori_count = parent_count[0]

    for i in range(len(parent_count) - 1):
        now_ind, now_count = parent_count[i + 1]
        if len(cost_heap[ori_ind]) == 0:
            print('Impossible')
            sys.exit()
        ori_cost, o_ind = heapq.heappop(cost_heap[ori_ind])
        now_cost, n_ind = heapq.heappop(cost_heap[now_ind])
        ans += ori_cost + now_cost

        if len(cost_heap[n_ind]) > 0:
            a, b = heapq.heappop(cost_heap[n_ind])
            heapq.heappush(cost_heap[ori_ind], (a, b))
        if len(cost_heap[o_ind]) > 0 and o_ind != ori_ind:
            a, b = heapq.heappop(cost_heap[o_ind])
            heapq.heappush(cost_heap[ori_ind], (a, b))

    print(ans)



if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
from collections import Counter, defaultdict
import sys
sys.setrecursionlimit(10 ** 5 + 10)
input = sys.stdin.readline
from math import factorial
import heapq, bisect
import math
import itertools


import queue
from collections import deque


parent_data = []

def union(ind):
    global parent_data
    if parent_data[ind] != ind:
        parent_data[ind] = union(parent_data[ind])
    return parent_data[ind]


def main():
    node_num, edge_num = map(int, input().split())
    cost_data = list(map(int, input().split()))
    data = [list(map(int, input().split())) for i in range(edge_num)]

    global parent_data
    parent_data = [i for i in range(node_num)]

    for i in range(edge_num):
        a, b = data[i]
        parent_a = union(a)
        parent_b = union(b)
        parent_data[min(parent_a, parent_b)] = max(parent_a, parent_b)

    parent_count = defaultdict(int)
    cost_heap = defaultdict(list)
    for i in range(node_num):
        parent_i = union(i)
        parent_count[parent_i] += 1
        heapq.heappush(cost_heap[parent_i], cost_data[i])

    parent_count = list(sorted(parent_count.items(), key=lambda x:-x[1]))

    ans = 0
    ori_ind, ori_count = parent_count[0]
    for i in range(len(parent_count) - 1):
        now_ind, now_count = parent_count[i + 1]
        if len(cost_heap[ori_ind]) == 0:
            print('Impossible')
            sys.exit()
        ori_cost = heapq.heappop(cost_heap[ori_ind])
        now_cost = heapq.heappop(cost_heap[now_ind])
        ans += ori_cost + now_cost
        cost_heap[ori_ind] = heapq.merge(cost_heap[ori_ind], cost_heap[now_ind])

    print(ans)





if __name__ == '__main__':
    main()