diff --git a/data_structures/linked_list/__init__.py b/data_structures/linked_list/__init__.py index a050adba4..3ddfea5c5 100644 --- a/data_structures/linked_list/__init__.py +++ b/data_structures/linked_list/__init__.py @@ -7,9 +7,11 @@ class Node: class LinkedList: def __init__(self): self.head = None + self.size = 0 def add(self, item): self.head = Node(item, self.head) + self.size += 1 def remove(self): if self.is_empty(): @@ -17,7 +19,28 @@ class LinkedList: else: item = self.head.item self.head = self.head.next + self.size -= 1 return item def is_empty(self): return self.head is None + + def __len__(self): + """ + >>> linked_list = LinkedList() + >>> len(linked_list) + 0 + >>> linked_list.add("a") + >>> len(linked_list) + 1 + >>> linked_list.add("b") + >>> len(linked_list) + 2 + >>> _ = linked_list.remove() + >>> len(linked_list) + 1 + >>> _ = linked_list.remove() + >>> len(linked_list) + 0 + """ + return self.size