hello-algo/codes/python/chapter_backtracking/backtrack_find_constrained_paths.py

"""
File: backtrack_find_constrained_path.py
Created Time: 2023-04-15
Author: Krahets (krahets@163.com)
"""

import sys, os.path as osp

sys.path.append(osp.dirname(osp.dirname(osp.abspath(__file__))))
from modules import *


def is_solution(state: list[TreeNode]) -> bool:
    """判断当前状态是否为解"""
    return state and state[-1].val == 7


def record_solution(state: list[TreeNode], res: list[list[TreeNode]]):
    """记录解"""
    res.append(list(state))


def is_valid(state: list[TreeNode], choice: TreeNode) -> bool:
    """判断在当前状态下，该选择是否合法"""
    return choice is not None and choice.val != 3


def make_choice(state: list[TreeNode], choice: TreeNode):
    """更新状态"""
    state.append(choice)


def undo_choice(state: list[TreeNode], choice: TreeNode):
    """恢复状态"""
    state.pop()


def backtrack(state: list[TreeNode], choices: list[TreeNode], res: list[list[TreeNode]]):
    """回溯算法"""
    # 检查是否为解
    if is_solution(state):
        # 记录解
        record_solution(state, res)
        return
    # 遍历所有选择
    for choice in choices:
        # 剪枝：检查选择是否合法
        if is_valid(state, choice):
            # 尝试：做出选择，更新状态
            make_choice(state, choice)
            backtrack(state, [choice.left, choice.right], res)
            # 回退：撤销选择，恢复到之前的状态
            undo_choice(state, choice)


"""Driver Code"""
if __name__ == "__main__":
    root = list_to_tree([1, 7, 3, 4, 5, 6, 7])
    print("\n初始化二叉树")
    print_tree(root)
    
    # 回溯算法
    res = []
    backtrack(state=[], choices=[root], res=res)

    print("\n输出所有根节点到节点 7 的路径，要求路径中不包含值为 3 的节点")
    for path in res:
        print([node.val for node in path])