526互联

用颜色标记法，实现树的前中后序遍历

发布时间 2023-06-30 19:46:08作者: w一蓑烟雨任平生w

使用颜色标记法，实现树的前中后序遍历

package algorithm;

import java.util.*;
import java.util.function.BiConsumer;

/**
 * 树的前中后序遍历 - 颜色标记法
 */
public class TreeTraversal {
    // white - 未访问过的节点
    private static final int WHITE = 0;
    // red - 访问过的节点
    private static final int RED = 0;
    private static final EnumMap<TraversalStrategy, BiConsumer<ColoredTreeNode, Stack<ColoredTreeNode>>> strategyHandler = new EnumMap<>(TraversalStrategy.class);

    public List<Integer> traversal(TreeNode root, TraversalStrategy strategy) {
        if (root == null) {
            return new ArrayList<>();
        }

        BiConsumer<ColoredTreeNode, Stack<ColoredTreeNode>> strategyConsumer = Optional.ofNullable(strategyHandler.get(strategy))
                .orElseThrow(() -> new RuntimeException("Unknown strategy: " + strategy));
        List<Integer> result = new ArrayList<>();
        Stack<ColoredTreeNode> stack = new Stack<>();
        stack.push(new ColoredTreeNode(WHITE, root));

        while (!stack.isEmpty()) {
            ColoredTreeNode curr = stack.pop();

            if (curr.color == WHITE) { // 如果节点未访问过
                // 使用对应策略，将根节点和左、右节点推入栈
                strategyConsumer.accept(curr, stack);
            } else {
                result.add(curr.treeNode.val);
            }
        }
        return result;
    }

    // 中序遍历
    BiConsumer<ColoredTreeNode, Stack<ColoredTreeNode>> inorderTraversal = (coloredTreeNode, stack) -> {
        // 中序遍历的顺序是: 左 -> 根 -> 右
        // 而栈是先进后出的结构，所以向栈中推数的顺序刚好相反，右 -> 根 -> 左
        if (coloredTreeNode.treeNode.right != null) {
            stack.push(new ColoredTreeNode(WHITE, coloredTreeNode.treeNode.right));
        }
        // curr节点已访问过，因此将颜色改为红色，再次推入栈中
        stack.push(coloredTreeNode.setColor(RED));
        if (coloredTreeNode.treeNode.left != null) {
            stack.push(new ColoredTreeNode(WHITE, coloredTreeNode.treeNode.left));
        }
    };

    // 前序遍历
    BiConsumer<ColoredTreeNode, Stack<ColoredTreeNode>> preorderTraversal = (coloredTreeNode, stack) -> {
        // 前序遍历的顺序是: 根 -> 左 -> 右
        // 而栈是先进后出的结构，所以向栈中推数的顺序刚好相反，右 -> 左 -> 根
        if (coloredTreeNode.treeNode.right != null) {
            stack.push(new ColoredTreeNode(WHITE, coloredTreeNode.treeNode.right));
        }
        if (coloredTreeNode.treeNode.left != null) {
            stack.push(new ColoredTreeNode(WHITE, coloredTreeNode.treeNode.left));
        }
        stack.push(coloredTreeNode.setColor(RED));
    };

    // 后序遍历
    BiConsumer<ColoredTreeNode, Stack<ColoredTreeNode>> postorderTraversal = (coloredTreeNode, stack) -> {
        // 后序遍历的顺序是: 左 -> 右 -> 根
        // 因此入栈顺序相反，是 根 -> 右 -> 左
        stack.push(coloredTreeNode.setColor(RED));
        if (coloredTreeNode.treeNode.right != null) {
            stack.push(new ColoredTreeNode(WHITE, coloredTreeNode.treeNode.right));
        }
        if (coloredTreeNode.treeNode.left != null) {
            stack.push(new ColoredTreeNode(WHITE, coloredTreeNode.treeNode.left));
        }
    };

    {
        // 定义不同遍历策略需要使用的栈处理方式
        strategyHandler.put(TraversalStrategy.INORDER, inorderTraversal);
        strategyHandler.put(TraversalStrategy.PREORDER, preorderTraversal);
        strategyHandler.put(TraversalStrategy.POSTORDER, postorderTraversal);
    }

    public enum TraversalStrategy {
        INORDER, // 中序
        PREORDER, // 前序
        POSTORDER // 后序
    }

    public static class ColoredTreeNode {
        // 树节点的颜色
        int color;
        // 被包装的树节点
        TreeNode treeNode;

        public ColoredTreeNode(int color, TreeNode treeNode) {
            this.color = color;
            this.treeNode = treeNode;
        }

        public ColoredTreeNode setColor(int color) {
            this.color = color;
            return this;
        }
    }

    public static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        public TreeNode() {
        }

        public TreeNode(int val) {
            this.val = val;
        }

        public TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }
}