建立打印LeetCode中的二叉树

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

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

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

再增加建树的成员方法

    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?这个问答，再添加树的打印方法

    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();
    }

测试

使用以下测试函数

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);
    }
}

获得的测试结果：

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

完整的TreeNode类

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();
    }

}

Python版本

from typing import List


class ListNode:
    def __init__(self, val: int = 0, next_=None):
        self.val = val
        self.next = next_


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

    @staticmethod
    def creatTreeNode(nums: List[int]):
        """
        输入数组层序遍历构造二叉树
        [1, 2, 3, None, 4]
            1
           / \
          2   3
           \
           4
        """
        if len(nums) == 0:
            return None
        nums_copy = nums.copy()

        popped = nums_copy.pop(0)
        root = TreeNode(popped)
        q = [root]
        while len(nums_copy) > 0:
            node = q.pop(0)

            left = nums_copy.pop(0)
            if left is not None:
                node.left = TreeNode(left)
                q.append(node.left)
            if len(nums_copy) == 0:
                break
            right = nums_copy.pop(0)
            if right is not None:
                node.right = TreeNode(right)
                q.append(node.right)
        return root

    def __str_helper(self, prefix, isTail, sb):
        if self.right is not None:
            if isTail:
                sb = self.right.__str_helper("" + prefix + "│   ", False, sb)
            else:
                sb = self.right.__str_helper("" + prefix + "    ", False, sb)

        if isTail:
            sb = sb + prefix + "└── " + str(self.val) + "\n"
        else:
            sb = sb + prefix + "┌── " + str(self.val) + "\n"

        if self.left is not None:
            if isTail:
                sb = self.left.__str_helper("" + prefix + "    ", True, sb)
            else:
                sb = self.left.__str_helper("" + prefix + "│   ", True, sb)
        return sb

    def __str__(self):
        return self.__str_helper("", True, "")


if __name__ == '__main__':
    root = TreeNode.creatTreeNode([1, 2, 3, None, 5])
    print(root)