py-leetcode/1912.py at main · uchitsa/py-leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from sortedcontainers import SortedList
from collections import defaultdict

class MovieRentingSystem:
    def __init__(self, n: int, entries: List[List[int]]):
        self.prices = {}
        self.availy = defaultdict(SortedList)
        self.rented = SortedList()

        for shop, movie, price in entries:
            self.prices[(shop, movie)] = price
            self.availy[movie].add((price, shop))

    def search(self, movie: int) -> List[int]:
        if movie not in self.availy:
            return []

        result = []
        for price, shop in self.availy[movie][:5]:
            result.append(shop)
        return result

    def rent(self, shop: int, movie: int) -> None:
        price = self.prices[(shop, movie)]
        self.availy[movie].remove((price, shop))
        self.rented.add((price, shop, movie))

    def drop(self, shop: int, movie: int) -> None:
        price = self.prices[(shop, movie)]
        self.rented.remove((price, shop, movie))
        self.availy[movie].add((price, shop))

    def report(self) -> List[List[int]]:
        res = []
        for price, shop, movie in self.rented[:5]:
            res.append([shop, movie])
        return res