Algorithms #3 ver.2 - better

import random
import sys


class Heap():

    def __init__(self, array):
        self.array = []  # tuple (key, value)
        self.pos = {}  # position of key in array
        self.size = len(array)  # number of elements in min heap
        for i, item in enumerate(array):
            self.array.append((item[0], item[1]))
            self.pos[item[0]] = i
        for i in range(self.size // 2, -1, -1):
            self.heapify(i)

        # display items of heap
    def display(self):
        print('Min array =', end=' ')
        for i in range(self.size):
            print(f'({self.array[i][0]} : {self.array[i][1]})', end=' ')
        print()
    # is heap empty ?

    def isEmpty(self):
        return self.size == 0
    # i is an array index
    # modify array so that i roots a heap, down-heap

    def heapify(self, i):
        smallest = i
        le = 2 * i + 1
        ri = 2 * i + 2
        if le < self.size and self.array[le][1] < self.array[smallest][1]:
            smallest = le
        if ri < self.size and self.array[ri][1] < self.array[smallest][1]:
            smallest = ri
        if smallest != i:
            # update pos
            self.pos[self.array[smallest][0]] = i
            self.pos[self.array[i][0]] = smallest
            # swap
            self.array[smallest], self.array[i] = self.array[i], self.array[smallest]
            self.heapify(smallest)

    # return the min element of the heap
    def getMin(self):
        if self.size == 0:
            return None
        return self.array[0]

    # return and remove the min element of the heap
    def pop(self):
        if self.size == 0:
            return None
        root = self.array[0]
        lastNode = self.array[self.size - 1]
        self.array[0] = lastNode
        del self.array[self.size - 1]
        # update pos
        self.pos[lastNode[0]] = 0
        del self.pos[root[0]]
        self.size -= 1
        self.heapify(0)
        return root

    # item (key, value) to push
    # modify array to include item
    def push(self, item):
        # push an item at the end with big value
        if self.size < len(self.array):
            self.array[self.size] = (item[0], float('inf'))
        else:
            self.array.append((item[0], float('inf')))
        self.pos[item[0]] = self.size
        self.size += 1
        self.decreaseKey(item)

    # decrease value of item (key, value)
    # with value smaller than current
    def decreaseKey(self, item):
        i = self.pos[item[0]]
        val = item[1]
        # new value must be smaller than current
        if self.array[i][1] <= val:
            return
        self.array[i] = item
        # check if is smaller than parent
        p = (i - 1) // 2
        while i > 0 and self.array[i][1] < self.array[p][1]:
            # update pos
            self.pos[self.array[i][0]] = p
            self.pos[self.array[p][0]] = i
            # swap
            self.array[p], self.array[i] = self.array[i], self.array[p]
            i = p
            p = (i - 1) // 2

    # increase value of item (key, value),
    # with value greater than current
    def increaseKey(self, item):
        i = self.pos[item[0]]
        val = item[1]
        # new value must be greater than current
        if self.array[i][1] >= val:
            return None
        self.array[i] = item
        # check children
        self.heapify(i)

    # change value of item (key, value),
    # decreaseKey or increaseKey, Running Time: O(log n)
    def changeKey(self, item):
        i = self.pos[item[0]]
        val = item[1]
        if val < self.array[i][1]:
            self.decreaseKey(item)
        if val > self.array[i][1]:
            self.increaseKey(item)

    # check if exists an item with key = key
    def isInHeap(self, key):
        if self.pos.get(key) == None:
            return False
        if self.pos[key] < self.size:
            return True
        return False

    # get the value of item with key = key
    def getValueMinHeap(self, key):
        if self.pos.get(key) == None:
            return None
        if self.pos[key] < self.size:
            return self.array[self.pos[key]][1]
        return None

    # This function deletes key at index i. It first reduces
    # value to minus infinite and then calls pop()
    def deleteKey(self, item):
        self.decreaseKey((item[0], float('-inf')))
        self.pop()


class Problem():
    def __init__(self):
        self.maxHeap = Heap([])
        self.minHeap = Heap([])

    def addNum(self, i):
        if self.maxHeap.isInHeap(i):
            self.maxHeap.changeKey((i[0], i[1]))
            return
        elif self.minHeap.isInHeap(i):
            self.minHeap.changeKey((i[0], i[1]))
            return
        if not self.maxHeap.array or i[1] < -self.maxHeap.array[0][1]:
            self.maxHeap.push((i[0], -i[1]))
        else:
            self.minHeap.push(i)

    def balance(self):
        if self.minHeap.size - self.maxHeap.size >= 2:
            temp = self.minHeap.pop()
            self.maxHeap.push((temp[0], -temp[1]))
        elif self.maxHeap.size - self.minHeap.size >= 2:
            temp = self.maxHeap.pop()
            self.minHeap.push((temp[0], -temp[1]))

    def median(self):
        # print('Max heap:\t{}\nMin heap:\t{}'.format(
        #    self.maxHeap.array, self.minHeap.array))
        if self.minHeap.size == self.maxHeap.size:
            return ((self.minHeap.array[0][1] - self.maxHeap.array[0][1]) / 2)
        elif self.minHeap.size > self.maxHeap.size:
            return float(self.minHeap.array[0][1])
        else:
            return float(-self.maxHeap.array[0][1])


if __name__ == '__main__':
    prob = Problem()
    random.seed()
    for i in range(500000):
        # prob.addNum(random.randint(1, 1000))
        prob.addNum(((str(random.randint(0, 2000)) + ('N, ' if random.randint(0, 2) % 2 == 0 else 'S, ') +
                    str(random.randint(0, 2000)) + ('E' if random.randint(0, 2) % 2 == 0 else 'W')), random.randrange(-2000, 3000)/100))
        prob.balance()
    print("The median temperature is {} C.".format(prob.median()))