def test_size(self): node = BSTNode("e") assert node.size() == 1 node.add("g") node.add("o") node.add("h") node.add("k") node.add("m") node.add("l") assert node.size() == 7
class MovieLib: """A movie library. Implemented using a BST. """ def __init__(self): """Initialise a movie library.""" self.bst = None def __str__(self): """Return a string representation of the library. The string will be created by an in-order traversal. """ if self.bst is not None: return str(self.bst) def size(self): """Return the number of movies in the library.""" if self.bst is None: return 0 return self.bst.size() def search(self, title): """Return Movie with matching title if there, or None. Args: title: a string representing a movie title. """ if self.bst is not None: return self.bst.search(Movie(title)) def add(self, title, date, runtime): """Add a new movie to the library. Args: title - the title of the movie date - the date the movie was released runtime - the running time of the movie Returns: the movie file that was added, or None """ add_movie = Movie(title, date, runtime) if self.bst is None: self.bst = BSTNode(add_movie) return add_movie return self.bst.add(add_movie) def remove(self, title): """Remove and return the a movie object with the given title, if there. Args: title - the title of the movie to be removed """ if self.bst is not None: removed = self.bst.remove(Movie(title)) if self.bst.size() == 0: # If tree is now empty self.bst = None # Set reference to None return removed def _testadd(): library = MovieLib() library.add("Memento", "11/10/2000", 113) print(str(library)) print("> adding Melvin and Howard") library.add("Melvin and Howard", "19/09/1980", 95) print(str(library)) print("> adding a second version of Melvin and Howard") library.add("Melvin and Howard", "21/03/2007", 112) print(str(library)) print("> adding Mellow Mud") library.add("Mellow Mud", "21/09/2016", 92) print(str(library)) print("> adding Melody") library.add("Melody", "21/03/2007", 113) print(str(library)) return library def _test(): library = MovieLib() library.add("B", "b", 1) print("Library:", library) print("adding", "A") library.add("A", "a", 1) print("Library:", library) print("removing", "A") library.remove("A") print("Library:", library) print("adding", "C") library.add("C", "c", 1) print("Library:", library) print("removing", "C") library.remove("C") print("Library:", library) print("adding", "F") library.add("F", "f", 1) print("Library:", library) print("removing", "B") library.remove("B") print("Library:", library) print("adding", "C") library.add("C", "c", 1) print("Library:", library) print("adding", "D") library.add("D", "d", 1) print("Library:", library) print("adding", "C") library.add("C", "c", 1) print("Library:", library) print("adding", "E") library.add("E", "e", 1) print("Library:", library) print("removing", "B") library.remove("B") print("Library:", library) print("removing", "D") library.remove("D") print("Library:", library) print("removing", "C") library.remove("C") print("Library:", library) print("removing", "E") library.remove("E") print("Library:", library) print("adding", "L") library.add("L", "l", 1) print("Library:", library) print("adding", "H") library.add("H", "h", 1) print("Library:", library) print("adding", "I") library.add("I", "i", 1) print("Library:", library) print("adding", "G") library.add("G", "g", 1) print("Library:", library) print("removing", "L") library.remove("L") print("Library:", library) print("removing", "H") library.remove("H") print("Library:", library) print("removing", "I") library.remove("I") print("Library:", library) print("removing", "G") library.remove("G") print("Library:", library)