Design Movie Rental System LeetCode solution in Python

1
# Runtime: 4131 ms (Top 50.00%) | Memory: 98.7 MB (Top 51.92%)
2
import sortedcontainers as sc
3

4

5
class MovieRentingSystem:
6
    def __init__(self, n: int, entries: List[List[int]]):
7
        self.rented = sc.SortedList()  # Triples of price, shop, movie
8
        self.movies = {}  # Dictionary mapping movie -> sorted list of (price, shop)
9
        self.movie_and_shop_to_price = {}  # Dictionary mapping (movie, shop) -> price
10

11
        for shop, movie, price in entries:
12
            if movie not in self.movies:
13
                self.movies[movie] = sc.SortedList()
14
            self.movies[movie].add((price, shop))
15
            self.movie_and_shop_to_price[movie, shop] = price
16

17
    def search(self, movie: int) -> List[int]:
18
        if movie not in self.movies:
19
            return []
20
        return [x[1] for x in self.movies[movie][:5]]  # Returning the shop numbers
21

22
    def rent(self, shop: int, movie: int) -> None:
23
        my_price = self.movie_and_shop_to_price[movie, shop]
24
        self.rented.add((my_price, shop, movie))
25
        self.movies[movie].remove((my_price, shop))
26

27
    def drop(self, shop: int, movie: int) -> None:
28
        my_price = self.movie_and_shop_to_price[movie, shop]
29
        self.rented.remove((my_price, shop, movie))
30
        self.movies[movie].add((my_price, shop))
31

32
    def report(self) -> List[List[int]]:
33
        return [x[1:] for x in self.rented[:5]]  # Returning the (shop, movie) pairs