LeetCode 1912. Design Movie Rental System

Posted on Edited on In 
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
const int N = 10001;
struct Node {
    int shop, movie, price;
    bool operator < (const Node t) const {
        if (price != t.price) return price < t.price;
        if (shop != t.shop) return shop < t.shop;
        return movie < t.movie;
    }
};
set<Node> in[N]; // available.
set<Node> out; // rent.
map<pair<int, int>, int> pr; // <shop, movie> => price.
class MovieRentingSystem {
public:
    MovieRentingSystem(int n, vector<vector<int>>& entries) {
        for (int i = 0; i < N; ++i) in[i].clear();
        out.clear();
        pr.clear();
        for (auto& e : entries) {
            int shop = e[0], movie = e[1], price = e[2];
            in[movie].insert({shop, movie, price});
            pr[{shop, movie}] = price;
        }
    }
    
    vector<int> search(int movie) {
        vector<int> ans;
        for (auto& m : in[movie]) {
            ans.emplace_back(m.shop);
            if (ans.size() == 5) break;
        }
        return ans;
    }
    
    void rent(int shop, int movie) {
        auto it = in[movie].find({shop, movie, pr[{shop, movie}]});
        out.insert(*it);
        in[movie].erase(it);
    }
    
    void drop(int shop, int movie) {
        auto it = out.find({shop, movie, pr[{shop, movie}]});
        in[movie].insert(*it);
        out.erase(it);
    }
    
    vector<vector<int>> report() {
        vector<vector<int>> ans;
        for (auto& e : out) {
            ans.push_back({e.shop, e.movie});
            if (ans.size() == 5) break;
        }
        return ans;
    }
};
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
from collections import defaultdict
from sortedcontainers import SortedList

class MovieRentingSystem:

    def __init__(self, n: int, entries: List[List[int]]):
        self._prices = dict()
        self._available = defaultdict(SortedList)
        self._rent = SortedList()
        for shop, movie, price in entries:
            self._available[movie].add((price, shop))
            self._prices[(shop, movie)] = price

    def search(self, movie: int) -> List[int]:
        a = self._available
        if movie not in a: return []
        return [shop for (price, shop) in a[movie][:5]]

    def rent(self, shop: int, movie: int) -> None:
        price = self._prices[(shop, movie)]
        self._available[movie].discard((price, shop))
        self._rent.add((price, shop, movie))

    def drop(self, shop: int, movie: int) -> None:
        price = self._prices[(shop, movie)]
        self._rent.discard((price, shop, movie))
        self._available[movie].add((price, shop))

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