Метод size для MyTreeSet

lab2/src/main/java/ru/spbstu/telematics/java/MyTreeSet.java

@@ -14,9 +14,15 @@ public class MyTreeSet<E extends Comparable<E>> {
     }
 
     private Node root;
+    private int size;
 
     public MyTreeSet() {
         root = null;
+        size = 0;
+    }
+
+    public int size() {
+        return size;
     }
 
     public boolean add(E element) {
@@ -26,6 +32,7 @@ public class MyTreeSet<E extends Comparable<E>> {
 
         if (root == null) {
             root = new Node(element);
+            size++;
             return true;
         }
 
@@ -47,6 +54,7 @@ public class MyTreeSet<E extends Comparable<E>> {
             parentNode.right = newNode;
         }
 
+        size++;
         return true;
     }
 
@@ -72,6 +80,7 @@ public class MyTreeSet<E extends Comparable<E>> {
         }
         if (!contains(element)) return false;
         root = remove(root, element);
+        size--;
         return true;
     }
 

lab2/src/test/java/ru/spbstu/telematics/java/AppTest.java

@@ -58,4 +58,17 @@ class MyTreeSetTests {
         assertThrows(NullPointerException.class, () -> myTreeSet.remove(null));
         assertThrows(NullPointerException.class, () -> treeSet.remove(null));
     }
+
+    @Test
+    void testSize() {
+        assertEquals(myTreeSet.size(), treeSet.size());
+        treeSet.add(15);
+        myTreeSet.add(15);
+        assertEquals(myTreeSet.size(), treeSet.size());
+        treeSet.add(50);
+        myTreeSet.add(50);
+        assertEquals(myTreeSet.size(), treeSet.size());
+        assertEquals(myTreeSet.remove(50), treeSet.remove(50));
+        assertEquals(myTreeSet.size(), treeSet.size());
+    }
 }