控制台打印AVL树-java

【简介】本文展示如何在控制台打印一棵AVL树（即平衡树）。

节点结构

为举例方便，这里使用了最简单的结构，仅包含当前元素、左节点、右节点及高：

public class AvlNode<TreeNode> {

    TreeNode element;
    AvlNode<TreeNode> left;
    AvlNode<TreeNode> right;
    int height; // 叶节点高度为0，根节点高度最大

    public AvlNode(TreeNode element) {
        this(element, null, null);
    }

    public AvlNode(TreeNode element, AvlNode<TreeNode> left, AvlNode<TreeNode> right)         {
        this.element = element;
        this.left = left;
        this.right = right;
        this.height = 0;
    }

    @Override
    public String toString() {
        return "AvlNode{" +
                "element=" + element +
                ", left=" + (left == null ? null : left.element) +
                ", right=" + (right == null ? null : right.element) +
                ", height=" + height +
                '}';
        }
    }
}

public class TreeNode {

    protected int number;

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

    public int compareTo(@NotNull TreeNode node) {
        if (node == null) {
            throw new RuntimeException("The node can't be null!");
        }
        if (number > node.number) {
            return 1;
        } else if (number == node.number) {
            return 0;
        } else {
            return -1;
        }
    }

    @Override
    public String toString() {
        return "TreeNode{" +
                "number=" + number +
                '}';
    }
}

打印AVL树

【打印细则】只打印元素，不打印线条，元素为1~2位数字（不足2位补空格），左叶节点后面空2格（其父节点在上一行占2格），右叶节点后空4格（其父节点的父节点在往前2行，占居中2格）。如有需要，可稍作改造以打印更多的位数。

【打印原理】假设我们的AVL树是完美的理想AVL树，所有叶节点深度相同（对于深度不足的，在缺失的位置使用“空格+N”占位），于是从完美AVL树，可以准确计算得出各个元素前后需要空出的位置。
以下使用 org.junit 做单元测试打印，以及相关方法：

public class PrintTest {

    @Test
    public void print() {
        printTree(root);
    }

    public void printTree(AvlNode<TreeNode> root) {
        if (root == null) {
            System.out.println("N");
            return;
        }
        List<String> nodesList = new ArrayList<>();
        // 在控制台只能逐行打印，因此行数和树的高度相关
        for (int i = 0; i <= root.height; i++) {
            nodesList.add("");
        }
        // 将树的节点信息放入集合中
        pushNodesToList(nodesList, root, 0, true);
        // 逐行打印
        for (int i = 0; i < nodesList.size(); i++) {
            System.out.println(nodesList.get(i));
        }
    }

    private void pushNodesToList(List<String> nodesList, AvlNode<TreeNode> base, int depth, boolean isLeft) {
        if (base == null) {
            // 当前节点为空，则该节点及其子节点均不存在，构造虚拟节点
            pushVirtualNodesToList(nodesList, depth, isLeft);
            return;
        }
        // 不使用真实height，确保和完美AVL树一致
        int height = nodesList.size() - depth - 1;
        if (height < 0) {
            return;
        }
        String lineStr = nodesList.get(depth) 
                + getSpaces(beforeSpaces(height)) 
                + getNumberStr(base) 
                + getSpaces(afterSpaces(height, isLeft));
        nodesList.set(depth, lineStr);
        depth++;
        pushNodesToList(nodesList, base.left, depth, true);
        pushNodesToList(nodesList, base.right, depth, false);
    }

    private void pushVirtualNodesToList(List<String> nodesList, int depth, boolean isLeft) {
        int height = nodesList.size() - depth - 1;
        if (height < 0) {
            return;
        }
        String lineStr = nodesList.get(depth)
                +  getSpaces(beforeSpaces(height))
                + " N"
                + getSpaces(afterSpaces(height, isLeft));
        nodesList.set(depth, lineStr);
        if (height > 0) {
            pushVirtualNodesToList(nodesList, depth + 1, true);
            pushVirtualNodesToList(nodesList, depth + 1, false);
        }
    }

    // 关键算法，计算元素之前的空格数
    private int beforeSpaces(int height) {
        if (height == 0) {
            return 0;
        }
        if (height == 1) {
            return 2;
        }
        if (height > 1) {
            return (int) Math.round(5 * Math.pow(2, height - 1) - 3);
        }
        throw new RuntimeException("height can't be negative");
    }

    // 关键算法，计算元素之后的空格数
    private int afterSpaces(int height, boolean isLeft) {
        if (height == 0) {
            return isLeft ? 2 : 4;
        }
        if (height == 1) {
            return 6;
        }
        if (height > 1) {
            return (int) Math.round(5 * Math.pow(2, height - 1) + 1);
        }
        throw new RuntimeException("height can't be negative");
    }

    // 元素不足两位则空格补位
    private String getNumberStr(AvlNode<TreeNode> avlNode) {
        if (avlNode == null) {
            return " N";
        }
        return (avlNode.element.number < 10 ? " " : "") + avlNode.element.number;
    }

    // 获取指定数量的空格数
    public static String getSpaces(int count) {
        char space = ' ';
        char[] spaces = new char[count];
        for (int i = 0; i < count; i++) {
            spaces[i] = space;
        }
        return new String(spaces);
    }
}