# 1912. Design Movie Rental System from collections import defaultdict from sortedcontainers import SortedList class MovieRentingSystem(object): def __init__(self, n, entries): """ :type n: int :type entries: List[List[int]] """ self.shops = defaultdict(SortedList) self.shop_movie = {} self.rented = SortedList() for s, m, p in entries: self.shops[m].add((p, s)) self.shop_movie[s, m] = p def search(self, movie): """ :type movie: int :rtype: List[int] """ return [i for _, i in self.shops[movie][:5]] def rent(self, shop, movie): """ :type shop: int :type movie: int :rtype: None """ price = self.shop_movie[shop, movie] self.shops[movie].remove((price, shop)) self.rented.add((price, shop, movie)) def drop(self, shop, movie): """ :type shop: int :type movie: int :rtype: None """ price = self.shop_movie[shop, movie] self.shops[movie].add((price, shop)) self.rented.remove((price, shop, movie)) def report(self): """ :rtype: List[List[int]] """ return [[i, j] for _, i, j in self.rented[:5]] if __name__ == "__main__": obj = MovieRentingSystem( 3, [[0, 1, 5], [0, 2, 6], [0, 3, 7], [1, 1, 4], [1, 2, 7], [2, 1, 5]] ) print(obj.search(1)) obj.rent(0, 1) obj.rent(1, 2) print(obj.report()) obj.drop(1, 2) print(obj.search(2))