LeetCode: 2034. Stock Price Fluctuation

 

1. Using heaps with lazy removal

Below code uses 'MyHeap' data structure, a heap with lazy removal. Lazy removal here means, if we have to remove an element, we handle it by using another heap as a trash can.

import heapq

class MyHeap:
    def __init__(self, type='min'):
        self.heap = []
        self.removals = []
        self.sign = 1 if type == 'min' else -1

    def insert(self, elem):
        heapq.heappush(self.heap, self.sign * elem)

    def __getitem__(self, idx):
        return self.sign * self.heap[idx]

    def __len__(self):
        return len(self.heap) - len(self.removals)

    def remove(self, elem):
        if elem == self.sign * self.heap[0]:
            heapq.heappop(self.heap)
            while self.removals and self.heap[0] == self.removals[0]:
                heapq.heappop(self.heap)
                heapq.heappop(self.removals)
        else:
            heapq.heappush(self.removals, self.sign * elem)


class StockPrice:
    
    def __init__(self):
        # data[timestamp] = price
        self.data = {}
        self.maxtimestamp = -1
        self.maxheap = MyHeap('max')
        self.minheap = MyHeap('min')

    def update(self, timestamp: int, price: int) -> None:
        if timestamp in self.data:
            wrong_price = self.data[timestamp]            
            self.maxheap.remove(wrong_price)
            self.minheap.remove(wrong_price)
        self.data[timestamp] = price
        self.maxtimestamp = max(self.maxtimestamp, timestamp)
        self.maxheap.insert(price)
        self.minheap.insert(price)
        
    def current(self) -> int:
        return self.data[self.maxtimestamp]

    def maximum(self) -> int:
        return self.maxheap[0]
    
    def minimum(self) -> int:
        return self.minheap[0]


# Your StockPrice object will be instantiated and called as such:
# obj = StockPrice()
# obj.update(timestamp,price)
# param_2 = obj.current()
# param_3 = obj.maximum()
# param_4 = obj.minimum()

 

2. Using dictionary to update heaps

Idea from: https://leetcode.com/problems/stock-price-fluctuation/discuss/1513293/Python-Clean-2-Heaps-Commented-Code

Python - Clean 2 Heaps - Commented Code - LeetCode Discuss

import heapq
class StockPrice:
    
    def __init__(self):
        self.data = {}
        self.maxtimestamp = -1
        self.maxheap = []
        self.minheap = []

    def update(self, timestamp: int, price: int) -> None:
        self.data[timestamp] = price
        self.maxtimestamp = max(self.maxtimestamp, timestamp)
        heapq.heappush(self.maxheap, (-price, timestamp))
        heapq.heappush(self.minheap, (price, timestamp))

    def current(self) -> int:
        return self.data[self.maxtimestamp]

    def maximum(self) -> int:
        price, timestamp = heapq.heappop(self.maxheap)
        while self.data[timestamp] != -price:
            price, timestamp = heapq.heappop(self.maxheap)
        heapq.heappush(self.maxheap, (price, timestamp))
        return -price

    def minimum(self) -> int:
        price, timestamp = heapq.heappop(self.minheap)
        while self.data[timestamp] != price:
            price, timestamp = heapq.heappop(self.minheap)
        heapq.heappush(self.minheap, (price, timestamp))
        return price

 

3. Using SortedDict

Idea from: https://leetcode.com/problems/stock-price-fluctuation/discuss/1513413/JavaC%2B%2BPython-Strightforward-Solutions

[Java/C++/Python] Strightforward Solutions - LeetCode Discuss

from sortedcontainers import SortedDict

class StockPrice:
    
    def __init__(self):
        self.prices = SortedDict()
        self.timestamps = SortedDict()

    def update(self, timestamp: int, price: int) -> None:
        if timestamp in self.timestamps:
            wrong_price = self.timestamps[timestamp]
            self.prices[wrong_price].remove(timestamp)
            if not self.prices[wrong_price]:
                self.prices.pop(wrong_price)
        if price not in self.prices:
            self.prices[price] = set()
        self.prices[price].add(timestamp)
        self.timestamps[timestamp] = price
        
    def current(self) -> int:
        return self.timestamps.peekitem(-1)[1]

    def maximum(self) -> int:
        return self.prices.peekitem(-1)[0]

    def minimum(self) -> int:
        return self.prices.peekitem(0)[0]