Cod sursa(job #2689438)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 20 decembrie 2020 20:05:41
Problema Heapuri Scor 40
Compilator py Status done
Runda Arhiva educationala Marime 3.3 kb
from __future__ import (division, print_function, 
                        absolute_import, unicode_literals)

def read_gen(filename):
    with open(filename, 'rt') as fin:
        for line in fin:
            for el in line.split():
                yield int(el)


class MinHeap:


    def __init__(self):
        self.elems = [0] # elements of the heap, elem[1] is top node
        self.ord = [0] # ord[i] = index in elems of the i-th element
        self.back = [0] # back[j] = index in ord of the j element from elems
    
    def insert(self, element):
        """Insert element into the heap."""
        self.elems.append(element)
        self.ord.append(len(self.elems) - 1)
        self.back.append(len(self.ord) - 1)
        self._percolate_up(len(self.elems) - 1)
    
    def delete_nth(self, n):
        """Delete n-th element in the order of insertion.
        
        It cannot be deleted directly because it will break
        the  vector reppresentation of the heap.
        """
        heap_index = self.ord[n]
        if heap_index != len(self.elems) - 1:
            self.swap(heap_index, len(self.elems) - 1)
            self.elems.pop()
            self.back.pop()
            self._percolate_up(heap_index)
            self._percolate_down(heap_index)
        else:
            self.elems.pop()
            self.back.pop()

    def get_min(self):
        """Return the minimum element."""
        return self.elems[1]

    def swap(self, fst, snd):
        """Swap elements by keeping the internal structure consistent."""
        self.elems[fst], self.elems[snd] = self.elems[snd], self.elems[fst]
        index_ord1, index_ord2 = self.back[fst], self.back[snd]
        self.ord[index_ord1], self.ord[index_ord2] = self.ord[index_ord2], self.ord[index_ord1]
        self.back[fst], self.back[snd] = self.back[snd], self.back[fst]

    def _percolate_down(self, index):
        """Percolates down (fitlers) the element from index."""
        while index < len(self.elems):
            min_son = -1
            if index * 2 < len(self.elems):
                min_son = index * 2
            if (index * 2 + 1 < len(self.elems) 
                and self.elems[index * 2 + 1] < self.elems[min_son]):
                min_son = index * 2 + 1 
            if min_son == -1 or self.elems[index] < self.elems[min_son]:
                break
            else:
                self.swap(min_son, index)
                index = min_son

    def _percolate_up(self, index):
        """Percolates up (fitlers) the element from index."""
        while index > 1 and self.elems[index] < self.elems[index // 2]:
            self.swap(index, index // 2)
            index = index // 2

if __name__ == '__main__':
    min_heap = MinHeap()
    it = read_gen(filename='heapuri.in')

    with open('heapuri.out', 'wt') as fout:
        m = next(it)
        for _ in range(m):
            op = next(it)
            assert (op in set([1, 2, 3])), 'Invalid operation'
            if op == 1:
                x = next(it)
                min_heap.insert(x)
            elif op == 2:
                nth = next(it)
                min_heap.delete_nth(nth)
            elif op == 3:
                min_el = min_heap.get_min()
                fout.write('{0}\n'.format(min_el))