from bisect import insort class Library(object): class _Reservation(object): def __init__(self, from_, to, name): self.from_ = from_ self.to = to # pylint: disable=invalid-name self.name = name def overlapping(self, res): return self.name == res.name and self.to >= res.from_ and res.to >= self.from_ def includes(self, date): return self.from_ <= date <= self.to def __lt__(self, other): return self.from_ < other.from_ def __init__(self): self._users = set() self._books = {} #maps name to count self._reservations = [] #Reservations sorted by from def add_user(self, name): if name in self._users: return False self._users.add(name) return True def add_book(self, name): self._books[name] = self._books.get(name, 0) + 1 def get_book_count(self, name): return self._books.get(name, 0) def reserve_book(self, user, book, date_from, date_to): book_count = self._books.get(book, 0) if user not in self._users or date_from > date_to or book_count == 0: return False desired_reservation = Library._Reservation(date_from, date_to, book) relevant_reservations = [res for res in self._reservations if desired_reservation.overlapping(res)] + [desired_reservation] #we check that if we add this reservation then for every reservation record that starts #between date_from and date_to no more than book_count books are reserved. for from_ in [res.from_ for res in relevant_reservations]: if desired_reservation.includes(from_): if sum([rec.includes(from_) for rec in relevant_reservations]) > book_count: return False insort(self._reservations, desired_reservation) #not for speed, still O(n) return True def books(self): for book in self._books: for _ in range(self._books[book]): yield book def reservations(self): return ((res.from_, res.to, res.name) for res in self._reservations) def at_most_books(max_number_of_books): def decorator(cls): class NewBookCollection(cls): # pylint: disable=too-few-public-methods def add_book(self, name): if self.get_book_count(name) >= max_number_of_books: return False cls.add_book(self, name) return True return NewBookCollection return decorator def one_book(bookcollection): return at_most_books(1)(bookcollection)