Cod sursa(job #2536036)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 1 februarie 2020 14:00:39
Problema Heapuri Scor 0
Compilator py Status done
Runda Arhiva educationala Marime 2.12 kb
def read_gen(fname):
    with open(fname, 'rt') as fin:
        for line in fin:
            for val in line.split():
                yield int(val)

def peek(h):
    return h[1]

def percolate_up(h, pos, elems, i):
    if i > 1:
        father = i // 2
        if elems[h[father]] > elems[h[i]]:
            h[i], h[father] = h[father], h[i]
            pos[h[i]], pos[h[father]] = pos[h[father]], pos[h[i]]
            percolate_up(h, pos, elems, father)

def percolate_down(h, pos, elems, i):

    son1 = i * 2
    son2 = son1 + 1
    change = 0
    if son1 in h and son2 not in h:
        change = son1
    elif son1 in h and son2 in h:
        if elems[h[i]] > elems[h[son1]] or elems[h[i]] > elems[h[son2]]:
            if elems[h[son1]] < elems[h[son2]]:
                change = son1
            else:
                change = son2
        else:
            change = 0
    if change > 0:
        h[change], h[i] = h[i], h[change]
        pos[h[change]], pos[h[i]] = pos[h[i]], pos[h[change]]
        percolate_down(h, pos, elems, change)
    
if __name__ == "__main__":
    it = read_gen('grader_test1.in')
    n = next(it)
    """
        h = indices from elems
        pos = map from index in elems to the position in h (inverse mapping of h)"""
    pos, elems, h = {}, {}, {}
    nh, ne = 0, 0

    with open('heapuri.out', 'wt') as fout:
        for _ in range(n):
            t = next(it)
            if t == 1 or t == 2:
                x = next(it)
            if t == 1:
                ne += 1
                nh += 1
                elems[ne] = x
                h[nh] = ne 
                pos[ne] = nh
                percolate_up(h, pos, elems, nh)
            elif t == 2:
                if pos[x] != nh:
                    h[pos[x]] = h[nh]
                    pos[h[nh]] = pos[x]
                    del h[nh]
                    nh -= 1
                    percolate_up(h, pos, elems, pos[x])
                    percolate_down(h, pos, elems, pos[x])
                    pos[x] = -1
                else:
                    del h[nh]
                    nh -= 1
            else:
                fout.write('{}\n'.format(elems[peek(h)]))