Problem p03440 (python)

Fixed submission:
from collections import deque, defaultdict
from heapq import heapify, heappop, merge, heapreplace


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

    def _root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self._root(self.table[x])
            return self.table[x]

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

    def union(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2


def solve(aa, uft):
    roots = defaultdict(list)
    for x, a in enumerate(aa):
        roots[uft._root(x)].append(a)
    if len(roots) == 1:
        return 0
    if n < (len(roots) - 1) * 2:
        return 'Impossible'
    ans = 0
    queue = []
    for l in roots.values():
        l.sort()
        ans += l[0]
        queue.append(l[1:])
    lm = n - m - 2
    if lm > 0:
        mm = merge(*queue)
        for _ in range(lm):
            ans += next(mm)
    return ans


n, m = map(int, input().split())
aa = list(map(int, input().split()))
uft = UnionFind(n)
for _ in range(m):
    uft.union(*map(int, input().split()))
print(solve(aa, uft))

Buggy submission (verdict: Runtime Error):
from bisect import insort
from collections import deque
from heapq import heapify, heappop, merge, heapreplace, heappush


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

    def _root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self._root(self.table[x])
            return self.table[x]

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

    def union(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2


def solve(aa, uft):
    roots = {x: [] for x, r in enumerate(uft.table) if r < 0}
    if n < (len(roots) - 1) * 2:
        return 'Impossible'
    for x, a in enumerate(aa):
        roots[uft._root(x)].append(a)
    queue = []
    for l in roots.values():
        l.sort()
        queue.append(deque(l))
    queue.sort()
    queue = deque(queue)
    ans = 0
    while len(queue) > 1:
        # print(queue)
        buf = []
        dq1 = queue.popleft()
        dq2 = queue.popleft()
        if len(dq1) == 1:
            while len(dq2) == 1 and len(queue) > 0:
                buf.append(dq2)
                dq2 = queue.popleft()
        ans += dq1.popleft() + dq2.popleft()
        insort(queue, deque(merge(dq1, dq2)))
        if buf:
            queue.extendleft(reversed(buf))
    return ans


n, m = map(int, input().split())
aa = list(map(int, input().split()))
uft = UnionFind(n)
for _ in range(m):
    uft.union(*map(int, input().split()))
print(solve(aa, uft))

Buggy submission (verdict: Runtime Error):
from bisect import insort
from collections import deque
from heapq import heapify, heappop, merge, heapreplace, heappush


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

    def _root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self._root(self.table[x])
            return self.table[x]

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

    def union(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2


def solve(aa, uft):
    roots = {x: [] for x, r in enumerate(uft.table) if r < 0}
    if n < (len(roots) - 1) * 2:
        return 'Impossible'
    for x, a in enumerate(aa):
        roots[uft._root(x)].append(a)
    queue = []
    for l in roots.values():
        l.sort()
        queue.append(deque(l))
    queue.sort()
    queue = deque(queue)
    ans = 0
    while len(queue) > 1:
        # print(queue)
        buf = []
        dq1 = queue.popleft()
        dq2 = queue.popleft()
        if len(dq1) == 1:
            while len(dq2) == 1 and len(queue) > 1:
                buf.append(dq2)
                dq2 = queue.popleft()
        ans += dq1.popleft() + dq2.popleft()
        insort(queue, deque(merge(dq1, dq2)))
        if buf:
            queue = deque(merge(queue, buf))
    return ans


n, m = map(int, input().split())
aa = list(map(int, input().split()))
uft = UnionFind(n)
for _ in range(m):
    uft.union(*map(int, input().split()))
print(solve(aa, uft))

Buggy submission (verdict: Runtime Error):
from bisect import insort
from collections import deque
from heapq import heapify, heappop, merge, heapreplace, heappush


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

    def _root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self._root(self.table[x])
            return self.table[x]

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

    def union(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2


def solve(aa, uft):
    roots = {x: [] for x, r in enumerate(uft.table) if r < 0}
    if n < (len(roots) - 1) * 2:
        return 'Impossible'
    for x, a in enumerate(aa):
        roots[uft._root(x)].append(a)
    queue = []
    for l in roots.values():
        l.sort()
        queue.append(deque(l))
    queue.sort()
    queue = deque(queue)
    ans = 0
    while len(queue) > 1:
        buf = []
        dq1 = queue.popleft()
        dq2 = queue.popleft()
        if len(dq1) == 1:
            while len(dq2) == 1 and len(queue) > 1:
                buf.append(dq2)
                dq2 = queue.popleft()
        ans += dq1.popleft() + dq2.popleft()
        insort(queue, deque(merge(dq1, dq2)))
        queue = deque(merge(queue, buf))
    return ans


n, m = map(int, input().split())
aa = list(map(int, input().split()))
uft = UnionFind(n)
for _ in range(m):
    uft.union(*map(int, input().split()))
print(solve(aa, uft))

Buggy submission (verdict: Runtime Error):
from collections import deque
from heapq import heapify, heappop, heappush, merge


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

    def _root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self._root(self.table[x])
            return self.table[x]

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

    def union(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2


def solve(aa, uft):
    roots = {x: [] for x, r in enumerate(uft.table) if r < 0}
    if n < (len(roots) - 1) * 2:
        return 'Impossible'
    for x, a in enumerate(aa):
        roots[uft._root(x)].append(a)
    queue = []
    for l in roots.values():
        l.sort()
        queue.append(deque(l))
    heapify(queue)
    ans = 0
    while len(queue) > 1:
        dq1 = heappop(queue)
        dq2 = heappop(queue)
        ans += dq1.popleft() + dq2.popleft()
        heappush(queue, deque(merge(dq1, dq2)))
    return ans


n, m = map(int, input().split())
aa = list(map(int, input().split()))
uft = UnionFind(n)
for _ in range(m):
    uft.union(*map(int, input().split()))
print(solve(aa, uft))

Buggy submission (verdict: Runtime Error):
from heapq import heapify, heappop, heappush


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

    def _root(self, x):
        if self.table[x] < 0:
            return x
        else:
            self.table[x] = self._root(self.table[x])
            return self.table[x]

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

    def union(self, x, y):
        r1 = self._root(x)
        r2 = self._root(y)
        if r1 == r2:
            return
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            if d1 == d2:
                self.table[r1] -= 1
        else:
            self.table[r1] = r2


def solve(aa, uft):
    roots = {x: [] for x, r in enumerate(uft.table) if r < 0}
    if n < (len(roots) - 1) * 2:
        return 'Impossible'
    for x, a in enumerate(aa):
        roots[uft._root(x)].append((a, x))
    queue = []
    for l in roots.values():
        heapify(l)
        queue.append((l[0][0], l))
    heapify(queue)
    ans = 0
    while len(queue) > 1:
        a1, q1 = heappop(queue)
        a2, q2 = heappop(queue)
        ans += a1 + a2
        heappop(q1)
        heappop(q2)
        q1 += q2
        heapify(q1)
        heappush(queue, (q1[0][0], q1))
    return ans


n, m = map(int, input().split())
aa = list(map(int, input().split()))
uft = UnionFind(n)
for _ in range(m):
    uft.union(*map(int, input().split()))
print(solve(aa, uft))