Iterating through a tree list and get unexpected results in the recursion

Question

I am developing an e-health system, which needs a tree-list structure to store information of the relationships between the diseases. I am trying to iterate through a tree structure looking like this:

Each node of this tree is a list that can potentially be any length, but the depth of this tree would be a limited constant. For each element in the list, it points at its child through another list. The basic structures of the root node and a general node look like:

class RootConditionNode:
    def __init__(self, diseaseName, leadTo=None, isTreated=False):
        self.diseaseName, self.leadTo, self.isTreated = diseaseName, leadTo, isTreated
        
    def __str__(self):
        return "Root condition: " + self.diseaseName + " - treated? " + str(self.isTreated)
    
    def __eq__(self, other):
        return self.diseaseName == other.diseaseName
     
class ConditionNode(RootConditionNode):
    
    def __init__(self, diseaseName, leadBy=None, leadTo=None, isTreated=False):
        """
        A general Node in the treelist that is not the root cause.
        DiseaseName can be modified to input SCTID later. 
        leadBy: what other disease is leading to the current disease, if this
            parameter is none, meaning this is the root cause of the whole disease tree
        leadTo: A list (or None) saying what this condition leads to
        """
        super().__init__(diseaseName, leadTo, isTreated)
        self.leadBy = leadBy
        
    def __str__(self):
        return "Condition: " + self.diseaseName + " - treated? " + str(self.isTreated)

The tree is built through:

class ConditionTree:
    
    def __init__(self, conditionDict):
        """
        Condition dictionary is a dictionary that each key is a condition,
        and the values include what this key condition is causing. 
        
        It's root is a list (at most of time only one element) that contains
        root cause(s) to each other conditions
        """
        self.root = []
        for key in conditionDict:
            self.insertToTree(key, conditionDict[key])
                
    def insertToTree(self, mainCond, condList):
        """
        Insert one key-value disease pair into the condition tree. 
        """
        
        leadTo = []
        
        if len(self.root) == 0:
            # First consider the first parent condition as the root condition. 
            rootNode = RootConditionNode(mainCond[0], isTreated=mainCond[1])
                        
            for cond in condList:
                leadTo.append(ConditionNode(cond[0], leadBy=rootNode, isTreated=cond[1]))
            rootNode.leadTo = leadTo
            self.root.append(rootNode)
            return
        
        parentNode = ConditionNode(mainCond[0], isTreated=mainCond[1])
        
        # If this parent node is already in the tree?
        existingNode = self.searchForNode(parentNode)
        if existingNode is not None:
            # Then just pop the lead to list to this node and finish insertion
            for cond in condList:
                newNode = ConditionNode(cond[0], leadBy=existingNode, isTreated=cond[1])
                leadTo.append(newNode)            
            existingNode.leadTo = leadTo
            return
        
        """   
        Running to here means the inserting parent node with its related diseases:
            a) Not inserting into an empty tree
            b) Not inserting into an existing node of the tree. 
    
        In this case, we first append this parent to the root and then managing its children.
        """
        parentNode = RootConditionNode(mainCond[0], isTreated=mainCond[1])
        self.root.append(parentNode)
        leadTo = []
        
        for cond in condList:
            childCond = ConditionNode(cond[0], leadBy=parentNode, isTreated=cond[1])
                
            """if this childCond is in root, needs to deprive its root level
            to the subchild level, but rebuilt the link. """
            if childCond in self.root:
                newNode = ConditionNode(cond[0], leadBy=parentNode, isTreated=cond[1])                
                for node in self.root:
                    if node == newNode:
                        childCond.leadTo = node.leadTo                
                self.root.remove(newNode)
                
            elif self.searchForNode(childCond, self.root):
                raise ValueError("Input disease dictionary has some logical errors!")
        
            
            # Append this child to the parent after all these testing.
            leadTo.append(childCond)
            
        
        parentNode.leadTo = leadTo
        # Test: print([ele.diseaseName for ele in parentNode.leadTo])

if __name__ == "__main__":
    
    conditionTree3 = ConditionTree({
        ("B", False) : [("E", False), ("F", True), ("G", False)], 
        ("C", False):[("J", True)],
        ("D", False) : [("H", True), ("I", False)], 
        ("A", False):[("B", False), ("C", False), ("D", False)],
        # ("H", True):[("J", True), ("K", True)],        
    })
    
    
    print("\n\nTest search:")
    result = conditionTree3.searchForNode(ConditionNode("H"))
    print("Search result: " + str(result))
    print()
    
    print("Condition tree 3: ")
    conditionTree3.printTree()
    #print("primary diagnose: "+ conditionTree3.getPrimaryDiagnoseCondition())
    print()

Issue occurs when I want to recursively iterate through the tree to find a node that matches, and this issue leads to inserting a deeper level node into the tree wrong (because insertToTree uses method searchForNode). The searchForNode looks like:

    def searchForNode(self, nodeToSearch, parent=None):
        
        if parent is None: # this means to traverse from the parent.
            parent = self.root
        
        for node in parent:
            if node.leadTo is not None:
                for each in node.leadTo:
                    print(each.diseaseName, end=' ')
                print()
            
            if node == nodeToSearch:
                nodeToReturn = ConditionNode(node.diseaseName, leadBy=node.leadBy, 
                    leadTo=node.leadTo, isTreated=node.isTreated)
                # TODO: Bug of this program here: when returned got nothing...
                print("Node to return: " + str(nodeToReturn))
                return nodeToReturn

            children = node.leadTo
            if children is not None:
                # Note you can't add return here, otherwise search will end 
                # somewhere halfway
                self.searchForNode(nodeToSearch, children)

While printing this tree leads to no issue at all:

    # Depth first traversing this tree will be similar to print this tree. 
    def printTree(self, parent=None, traverseLevel=0):
        # Used to just print out this tree
        
        if parent is None: # this means to traverse from the parent.
            parent = self.root
        
        for node in parent:        
            print("\t" * traverseLevel + str(node))
            children = node.leadTo
            if children is not None:
                self.printTree(children, traverseLevel + 1)

So the issue is, how do I implement my searchForNode, either through recursive iteration or other methods, so that the search returns a correct node object that I'm searching? Thanks.

Answer 1

In a code comment you write:

Note you can't add return here, otherwise search will end somewhere halfway

This is only true if you would do it unconditionally. But consider the case where the search is successful in the recursive search. Then you wouldn't know about it! You really need to check the result you get from this recursive call and distinguish between the two possible cases: either it was successful or not. Then deal with both cases accordingly: if it was successful, return that same result and exit the current function call. If it wasn't successful, continue the loop.

So replace this:

            self.searchForNode(nodeToSearch, children)

With this:

            result = self.searchForNode(nodeToSearch, children)
            if result:
                return result

The reason you don't have this issue in printTree is that there you want to visit the whole tree, no matter what, and you don't need the function to return something significant. Its only function is to output, not to return a value.