带特殊规则的多叉树DFS搜索问题及代码优化求助

需求规则 #

• 调用check_for_three_kind_of_response(question, node.value)返回结果分为三类：
  • -1：节点与问题无关，优先搜索当前节点的兄弟节点；若无兄弟节点，再搜索当前节点的子节点
  • 0：节点与问题相关但无法解答，直接搜索当前节点的子节点
  • 文本内容：直接返回该文本作为最终结果

原代码问题分析 #

原DFS函数的核心缺陷在于：当节点返回-1时直接返回-1，没有继续遍历当前节点的兄弟节点，导致漏掉同层级的潜在相关节点。例如在示例树中，若paragraph1返回-1，原代码会直接终止搜索，不会去遍历paragraph2节点。

修复后的代码 #

class TreeNode:
    def __init__(self, value, parent=None):
        self.value = value
        self.parent = parent
        self.children = []

def dfs(node, question):
    result = check_for_three_kind_of_response(question, node.value)
    
    # 直接返回有效文本结果
    if result not in ("-1", "0"):
        return result
    
    # 节点相关但无法解答，遍历子节点搜索
    if result == "0":
        for child in node.children:
            child_res = dfs(child, question)
            if child_res != "-1":
                return child_res
        return "-1"
    
    # 节点无关，优先遍历兄弟节点
    if result == "-1":
        # 遍历当前节点之后的兄弟节点
        if node.parent is not None:
            siblings = node.parent.children
            current_pos = siblings.index(node)
            for sibling in siblings[current_pos + 1:]:
                sibling_res = dfs(sibling, question)
                if sibling_res != "-1":
                    return sibling_res
        # 兄弟节点无结果，遍历当前节点的子节点
        for child in node.children:
            child_res = dfs(child, question)
            if child_res != "-1":
                return child_res
        return "-1"

修复逻辑说明 #

1. 文本结果直接返回：只要检测到有效回答文本，立即终止搜索并返回
2. 节点相关（返回0）：仅遍历当前节点的子节点，因为节点本身已确认与问题相关，无需再搜索兄弟
3. 节点无关（返回-1）：
   • 先从父节点的子节点列表中，找到当前节点之后的所有兄弟节点，逐个递归搜索
   • 若所有兄弟节点均无有效结果，再遍历当前节点的子节点进行搜索
   • 所有路径都无结果时，才返回-1

内容的提问来源于stack exchange，提问作者疲れた