finished some familytree functions

family_tree.py

@@ -2,12 +2,10 @@
 
 """
 TODO:
-Don't initialize with root.
-Find a way to get root afterwards.
 Find a way to get isolated nodes.
 
 current structure:
-    bin tree inside a dict
+    bin tree, nodes in a dict
 
 
 Problems to consider --
@@ -22,30 +20,24 @@ Problems taken care of(?) --
     Adding child to someone who already has two children.
 
 Questions:
-How does a dict compare object keys?
-How to deal with multiple trees?
 Is there only one root?
 
 Total nodes of the tree?
 11 generations = 2^11 - 1
 
-
-KEY = NAME????
-MAKE THE KEY ONLY THE PERSON'S NAME
-MAKE HELPER FUNC
-REWRITE EVERYTHING
-
-
-How do class variables work?
 """
 
 
 class FamilyTree:
     """A family tree.
 
+    Attributes --
+        family_tree: a dict containing the nodes of the familytree
+        root: the root of the tree
+
     Family Tree dict format --
         key = name of the person
-        value = a Person object
+        value = a Person object, acting as the node
     """
 
     def __init__(self):
@@ -56,19 +48,19 @@ class FamilyTree:
     def find_root(self):
         """Find the root of the tree.
 
-        note: work in progress
-        Start from arbitrary node,
-        then move upwards. (doesn't check for multiple roots)
-
-        or ?????
+        Start from arbitrary node, then move upwards.
 
-        go through every object in the dictionary
-        and check if it has a parent (slow)
+        Note: Will not work if there's more than one node without a parent
         """
-        return
+        arbitrary_node = next(iter(self.family_tree.values()))
+
+        while arbitrary_node.get_parent() is not None:
+            arbitrary_node = arbitrary_node.get_parent()
+
+        self.root = arbitrary_node
 
     def set_root(self, person):
-        """Set the root node."""
+        """Set the root node manually."""
         self.root = person.get_name()
 
     def insert_person(self, person):
@@ -91,23 +83,46 @@ class FamilyTree:
 
     def get_descendants(self, name):
         """Return a sorted list of a person's descendants."""
-        def get_preorder(person):
-            """Get the inorder traversal given a person as the root."""
+        def get_postorder(person):
+            """Get the postorder traversal given the root node."""
             res = []
             if person is not None:
                 lc, rc = person.get_children()
+                res += get_postorder(lc)
+                res += get_postorder(rc)
                 res.append(person.get_name())
-                res += get_preorder(lc)
-                res += get_preorder(rc)
             return res
 
-        descendants = get_preorder(self.family_tree[name])[1:]
+        descendants = get_postorder(self.family_tree[name])
+        descendants.pop()
         descendants.sort()
         return descendants
 
-    def __str__(self):
-        """Return the name."""
-        return self.name
+    def compute_generations(self):
+        """DFS to find and set the generation of each node."""
+        if self.root is None:
+            self.find_root()
+        # TODO
+        return
+
+    def slow_find_lca(self, a, b):
+        """Find the lca slowly.
+
+        two pointer method
+        """
+        return
+
+    def fast_find_lca(self, a, b):
+        """Find the lca fast.
+
+        Reduction to range min. query
+        """
+        return
+
+
+def dfs(A):
+    """Depth first search."""
+    pass
 
 
 class Person:
@@ -117,8 +132,9 @@ class Person:
         name: a string indicating the name
         age: an integer indicating the age
         sex: a boolean; male = True, female = False
-        children: a list of the children
-        parent: the parent
+        children: a list of the children nodes
+        parent: the parent node
+        generation: an integer indicating the generation (level)
     """
 
     def __init__(self, name, age, sex):
@@ -128,6 +144,7 @@ class Person:
         self.sex = sex
         self.children = [None, None]
         self.parent = None
+        self.generation = None
 
     def set_parent(self, parent):
         """Set parent.
@@ -157,6 +174,10 @@ class Person:
         """Return a list of the children."""
         return self.children
 
+    def get_name(self):
+        """Get name."""
+        return self.name
+
     def set_name(self, name):
         """Set name."""
         self.name = name
@@ -165,24 +186,33 @@ class Person:
         """Set age."""
         self.age = age
 
-    def set_sex(self, sex):
-        """Set sex (bool)."""
-        self.sex = sex
-
-    def get_name(self):
-        """Get name."""
-        return self.name
-
     def get_age(self):
         """Get age."""
         return self.age
 
+    def set_sex(self, sex):
+        """Set sex (bool)."""
+        self.sex = sex
+
     def get_sex(self):
         """Get sex (bool)."""
         return self.sex
 
+    def get_generation(self):
+        """Get generation."""
+        return self.generation
 
-if __name__ == "__main__":
+    def set_generation(self, generation):
+        """Set generation."""
+        self.generation = generation
+
+    def __str__(self):
+        """Return the name."""
+        return self.name
+
+
+def main():
+    """For testing."""
     ftree = FamilyTree()
     A = Person("A", 69, True)
     B = Person("B", 69, True)
@@ -192,6 +222,13 @@ if __name__ == "__main__":
     ftree.insert_person(A)
     ftree.insert_person(B)
     ftree.insert_person(C)
-    ftree.set_root(A)
-    print(ftree.get_ancestors(C.get_name()))
-    print(ftree.get_descendants(A.get_name()))
+    ftree.find_root()
+    ancestors = ftree.get_ancestors(C.get_name())
+    descendants = ftree.get_descendants(A.get_name())
+    print("The root of the tree is {}.".format(ftree.root))
+    print("ancestors of {}: {}".format("C", ancestors))
+    print("descendants of {}: {}".format("A", descendants))
+
+
+if __name__ == "__main__":
+    main()