aboutsummaryrefslogtreecommitdiff
path: root/design_movie_rental.py
blob: 384ed12eec01befcd31818c219a0e0d5b7759a71 (plain)
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
# 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))