Problem statement

https://leetcode.com/problems/design-movie-rental-system/

Solution

This is a problem with quite heavy statement, but in fact it can be solved pretty strarightforward using SortedList library in python. The idea is to keep several pieces of information:

  1. self.shops is dictionary: for every movie we keep sorted list of tuples (price, shop).
  2. self.shop_movie is dictionary, which allows us to answer what is price, given shop and movie.
  3. self.rented is SortedList of rented movies so far.

Now, lwt us go through our functions and check how they work:

  1. _ _init_ _: here we create empty objects, and then fill self.shops and self.shop_movie. We keep self.rented empty, because we did not rent any movies yet.
  2. search(movie): we go through first 5 elements of self.shops[movie] and return prices.
  3. rent(shop, movie): first, we find price, using our self.shop_movie dictionary and then remove this movie from given shop and add it to rented.
  4. drop(shop, movie): very similar to previous one, now we add movie back to shop and remove it from rented.
  5. report Here we look at self.rented and take 5 cheapes movies, that is 5 first elements.

Complexity

  1. _ _ init _ _: time is O(n log n), space is O(n).
  2. search(movie): time is O(log n), space is O(1).
  3. rent(shop, movie): time is O(log n), additional space is O(1).
  4. drop(shop, movie): time is O(log n), additional space is O(1).
  5. report: time is O(log n), space is O(1).

Code

from sortedcontainers import SortedList

class MovieRentingSystem:
    def __init__(self, n, entries):
        self.shops = defaultdict(SortedList)   #movie -> (price, shop)
        self.shop_movie = {}    #(shop, movie) -> price
        self.rented = SortedList()  # (price, shop, movie)
        for s, m, p in entries:
            self.shops[m].add((p, s))
            self.shop_movie[s, m] = p

    def search(self, movie):
        return [y for _,y in self.shops[movie][:5]]
        
    def rent(self, shop, movie):
        price = self.shop_movie[shop, movie]
        self.shops[movie].remove((price, shop))
        self.rented.add((price, shop, movie))

    def drop(self, shop, movie):
        price = self.shop_movie[shop, movie]
        self.shops[movie].add((price, shop))
        self.rented.remove((price, shop, movie))
        
    def report(self):
        return [[y,z] for _,y,z in self.rented[:5]]