建立打印力扣(Leetcode)中的二叉树-java

发表于 更新于 分类于 Changyan：

建立打印LeetCode中的二叉树

在刷力扣题时遇到二叉树的问题时，建树较为麻烦，打印树更加是不方便。

因为力扣中二叉树定义多为

1
2
3
4
5
6
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int val) { this.val = val; }
}

再增加建树的成员方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
public static TreeNode creatABitTree(Integer[] nums) {

    if (nums.length == 0 || nums[0] == null) return null;
    if (nums.length == 1) return new TreeNode(nums[0]);

    Queue<TreeNode> queue = new LinkedList<>();  // 得用队列
    // 根结点的处理
    TreeNode node = new TreeNode(nums[0]);
    queue.offer(node);

    TreeNode temp;
    for (int i = 1; i < nums.length; ) {
        temp = queue.poll();
        if (nums[i] != null) {
            assert temp != null;
            temp.left = new TreeNode(nums[i++]);
            queue.offer(temp.left);
        } else
            i++;
        if (i >= nums.length)
            break;
        if (nums[i] != null) {
            assert temp != null;
            temp.right = new TreeNode(nums[i++]);
            queue.offer(temp.right);
        } else
            i++;
    }


    return node;
}

参考 StackOverflow中How to print binary tree diagram?这个问答，再添加树的打印方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
public StringBuilder toString(StringBuilder prefix, boolean isTail, StringBuilder sb) {
    if (right != null) {
        right.toString(new StringBuilder().append(prefix).append(isTail ? "│   " : "    "), false, sb);
    }
    sb.append(prefix).append(isTail ? "└── " : "┌── ").append(val).append("\n");
    if (left != null) {
        left.toString(new StringBuilder().append(prefix).append(isTail ? "    " : "│   "), true, sb);
    }
    return sb;
}

@Override
public String toString() {
    return this.toString(new StringBuilder(), true, new StringBuilder()).toString();
}

测试

使用以下测试函数

1
2
3
4
5
6
7
8
class Main {
    public static void main(String[] args) {
        TreeNode treeNode = TreeNode.creatABitTree(new Integer[]{
                0,1,2,3,4,5,6
        });
        System.out.println(treeNode);
    }
}

获得的测试结果：

1
2
3
4
5
6
7
│       ┌── 6
│   ┌── 2
│   │   └── 5
└── 0
    │   ┌── 4
    └── 1
        └── 3

完整的TreeNode类

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import java.util.LinkedList;
import java.util.Queue;

/**
 * create time: 2021/3/9 上午 10:01
 *
 * @author DownUpZ
 */
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

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

    /***
     * 不选择 int[] 原因是 int 无法表示空
     * @param nums
     * @return
     */
    public static TreeNode creatABitTree(Integer[] nums) {

        if (nums.length == 0 || nums[0] == null) return null;
        if (nums.length == 1) return new TreeNode(nums[0]);

        Queue<TreeNode> queue = new LinkedList<>();  // 得用队列
        // 根结点的处理
        TreeNode node = new TreeNode(nums[0]);
        queue.offer(node);

        TreeNode temp;
        for (int i = 1; i < nums.length; ) {
            temp = queue.poll();
            if (nums[i] != null) {
                assert temp != null;
                temp.left = new TreeNode(nums[i++]);
                queue.offer(temp.left);
            } else
                i++;
            if (i >= nums.length)
                break;
            if (nums[i] != null) {
                assert temp != null;
                temp.right = new TreeNode(nums[i++]);
                queue.offer(temp.right);
            } else
                i++;
        }


        return node;
    }

    public StringBuilder toString(StringBuilder prefix, boolean isTail, StringBuilder sb) {
        if (right != null) {
            right.toString(new StringBuilder().append(prefix).append(isTail ? "│   " : "    "), false, sb);
        }
        sb.append(prefix).append(isTail ? "└── " : "┌── ").append(val).append("\n");
        if (left != null) {
            left.toString(new StringBuilder().append(prefix).append(isTail ? "    " : "│   "), true, sb);
        }
        return sb;
    }

    @Override
    public String toString() {
        return this.toString(new StringBuilder(), true, new StringBuilder()).toString();
    }

}