WIP for https://leetcode.com/problems/find-median-from-data-stream/

FindMedianFromDataStream.py

# https://leetcode.com/problems/find-median-from-data-stream/
# Work in progress
class MedianFinder:
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.number_count: Dict[int, int] = {}
        self.size = 0
        self.sorted_keys = None
        # self.c = 0

    def addNum(self, num: int) -> None:
        if num not in self.number_count:
            self.sorted_keys = None
        value = self.number_count.get(num, 0) + 1
        self.number_count[num] = value
        self.size += 1
        # print(f"addNum({num}) - {self.number_count}")

    def findMedian(self) -> float:
        if self.size == 1:
            for num in self.number_count.keys():
                return num

        # self.c += 1
        counted = 0
        last = None
        
        if not self.sorted_keys:
            self.sorted_keys = sorted(self.number_count.keys())

        for x in self.sorted_keys:
            count = self.number_count[x]
            # print(f"size: {self.size}, counted: {counted}, last: {last}, x: {x}, count: {count}")
            if counted + count > self.size / 2:
                if count == 1  or counted + 1 > self.size / 2:
                    if self.size % 2 == 0:
                        return (last + x) / 2
                    else:
                        return x
                
                return x
                
            last = x
            counted += count


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()