dfs function under construction

family_tree.py

@@ -98,12 +98,54 @@ class FamilyTree:
         descendants.sort()
         return descendants
 
-    def compute_generations(self):
-        """DFS to find and set the generation of each node."""
+    def precompute(self):
+        """Do the precomputations required to find LCA and degree.
+
+        DFS from root.
+        Along the way, set the generations (levels)
+        and fill up the arrays for the lca query (TODO)
+
+        Raises exception if a root cannot be found.
+        """
+        def set_gen(node):
+            """Set the generation of a node to be one more than its parent."""
+            if node.parent is not None:
+                node.set_generation(node.parent.generation+1)
+            else:
+                node.set_generation(1)
+
+        def pre_lca(node):
+            # TODO
+            self.dfs_order.append(node.get_name())
+            self.dfs_gen.append(node.get_generation())
+
         if self.root is None:
             self.find_root()
-        # TODO
-        return
+        if self.root is None:
+            raise Exception("Tree is empty.")
+        self.dfs(self.root.get_name(), set_gen)
+
+    def dfs(self, s, *funcs):
+        """Depth first search from s, doing the given functions on each node.
+
+        S = name of person in starting node.
+        *funcs = the operations to be done on the node
+        """
+        visited = []
+        current_node = self.family_tree[s]
+        for f in funcs:
+            f(current_node)
+        children = current_node.get_children()
+        lc = children[0]
+        rc = children[1]
+        if lc is not None:
+            if lc.get_name() not in visited:
+                visited.append(lc.get_name())
+                self.dfs(lc.get_name(), *funcs)
+        if rc is not None:
+            if rc.get_name() not in visited:
+                visited.append(rc.get_name())
+                self.dfs(rc.get_name(), *funcs)
 
     def slow_find_lca(self, a, b):
         """Find the lca slowly.
@@ -112,16 +154,23 @@ class FamilyTree:
         """
         return
 
-    def fast_find_lca(self, a, b):
+    def lca(self, a, b):
         """Find the lca fast.
 
         Reduction to range min. query
         """
         return
 
+    def degree(self, a, b):
+        """Find the degree between two cousins."""
+        pass
+
+    def relationship(self, p1, p2):
+        """Determine the relationship between two distinct people in the tree.
 
-def dfs(A):
-    """Depth first search."""
+        Input: Two people p1 and p2, where p1 != p2
+        Output: The relationship between p1 and p2, from the perspective of p1
+        """
         pass
 
 
@@ -217,17 +266,23 @@ def main():
     A = Person("A", 69, True)
     B = Person("B", 69, True)
     C = Person("C", 69, True)
-    B.set_parent(A)
-    C.set_parent(A)
+    D = Person("D", 69, True)
     ftree.insert_person(A)
     ftree.insert_person(B)
     ftree.insert_person(C)
+    ftree.insert_person(D)
+    B.set_parent(A)
+    C.set_parent(A)
+    D.set_parent(B)
     ftree.find_root()
-    ancestors = ftree.get_ancestors(C.get_name())
+    ancestors = ftree.get_ancestors(D.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("ancestors of {}: {}".format("D", ancestors))
     print("descendants of {}: {}".format("A", descendants))
+    ftree.precompute()
+    for key, node in ftree.family_tree.items():
+        print("{}: {}".format(node, node.get_generation()))
 
 
 if __name__ == "__main__":