AB3: BodyForceTreeMap working

src/BodyForceTreeMap.java

@@ -3,8 +3,8 @@
  * The number of key-value pairs is not limited.
  */
 public class BodyForceTreeMap {
-
-    //TODO: declare variables.
+    private int size = 0;
+    private BodyForceTreeMapItem root = null;
 
     /**
      * Adds a new key-value association to this map. If the key already exists in this map,
@@ -12,7 +12,37 @@ public class BodyForceTreeMap {
      * Precondition: key != null.
      */
     public Vector3 put(Body key, Vector3 value) {
-        //TODO: implement method.
+        if (root == null) {
+            root = new BodyForceTreeMapItem(key, value);
+            size++;
+            return null;
+        }
+
+        BodyForceTreeMapItem item = root;
+        while (item != null) {
+            if (item.key() == key) {
+                Vector3 old = item.value();
+                item.setValue(value);
+                return old;
+            } else if (item.key().mass() > key.mass()) {
+                if (item.left() != null) {
+                    item = item.left();
+                } else {
+                    item.setLeft(new BodyForceTreeMapItem(key, value));
+                    size++;
+                    break;
+                }
+            } else {
+                if (item.right() != null) {
+                    item = item.right();
+                } else{
+                    item.setRight(new BodyForceTreeMapItem(key, value));
+                    size++;
+                    break;
+                }
+            }
+        }
+
         return null;
     }
 
@@ -22,7 +52,16 @@ public class BodyForceTreeMap {
      * Precondition: key != null.
      */
     public Vector3 get(Body key) {
-        //TODO: implement method.
+        BodyForceTreeMapItem item = root;
+        while (item != null) {
+            if (item.key() == key) {
+                return item.value();
+            } else if (item.key().mass() > key.mass()) {
+                item = item.left();
+            } else {
+                item = item.right();
+            }
+        }
         return null;
     }
 
@@ -30,16 +69,38 @@ public class BodyForceTreeMap {
      * Returns 'true' if this map contains a mapping for the specified key.
      */
     public boolean containsKey(Body key) {
-        //TODO: implement method.
+        BodyForceTreeMapItem item = root;
+        while (item != null) {
+            if (item.key() == key) {
+                return true;
+            } else if (item.key().mass() > key.mass()) {
+                item = item.left();
+            } else {
+                item = item.right();
+            }
+        }
         return false;
     }
 
+    public int size() {
+        return this.size;
+    }
+
+    private String toString(BodyForceTreeMapItem item, String indent) {
+        if (item == null) {
+            return "";
+        }
+        String s = String.format("%s{%s: %s}\n", indent, item.key(), item.value());
+        s += this.toString(item.left(), indent + "| ");
+        s += this.toString(item.right(), indent + "| ");
+        return s;
+    }
+
     /**
      * Returns a readable representation of this map, in which key-value pairs are ordered
      * descending according to the mass of the bodies.
      */
     public String toString() {
-        //TODO: implement method.
-        return null;
+        return toString(root, "");
     }
 }