java实现二叉搜索树

前言

二叉搜索树(又叫二叉查找树，二叉排序树)是一种特殊的二叉树，根节点的值大于左孩子的值，小于右孩子的值。示例图如下

二叉树定义

/**
 * 自己实现二叉搜索树
 */
public class BST<E extends Comparable<E>> {

  /**
   * 根节点
   */
  private Node<E> root;
  /**
   * 树的节点数量
   */
  private int size;

  /**
   * 查询容量
   */
  public int size() {
    return size;
  }

  /**
   * 是否为空
   */
  public boolean isEmpty() {
    return size == 0;
  }

  private static class Node<E> {

    /**
     * 节点值
     */
    E data;
    /**
     * 左孩子
     */
    Node<E> left;
    /**
     * 右孩子
     */
    Node<E> right;
    /**
     * 父节点
     */
    Node<E> parent;

    Node(E data) {
      this.data = data;
    }

    @Override
    public String toString() {
      return String.valueOf(data);
    }
  }
}

树中存储的值必须是可比较的，这样才能排序。节点中的parent也可以没有。

添加节点

添加节点就是创建二叉搜索树的过程，不支持重复元素

递归实现

  /**
   * 添加元素
   */
  public void add(E e) {
    root = add(root, e);
    size++;
  }

  private Node<E> add(Node<E> root, E e) {
    if (root != null) {
      if (e.compareTo(root.data) > 0) {
        root.right = add(root.right, e);
        root.right.parent = root;
      } else if (e.compareTo(root.data) < 0) {
        root.left = add(root.left, e);
        root.left.parent = root;
      }
      return root;
    } else {
      return new Node<>(e);
    }
  }

非递归实现

  /**
   * 添加元素
   */
  public void addNoRecursion(E e) {
    //从根节点开始查找
    Node<E> cur = root;
    //待插入节点的父节点
    Node<E> parent = null;
    int cmp = 0;
    while (cur != null) {
      cmp = e.compareTo(cur.data);
      if (cmp > 0) {
        parent = cur;
        cur = cur.right;
      } else if (cmp < 0) {
        parent = cur;
        cur = cur.left;
      } else {
        return;
      }
    }
    Node<E> node = new Node<>(e);
    if (parent == null) {
      root = node;
    } else {
      if (cmp < 0) {
        parent.left = node;
      } else {
        parent.right = node;
      }
      node.parent = parent;
    }
    size++;
  } 

我们先创建一个二叉树，以便接下来的遍历和删除

public class Main {

  public static void main(String[] args) {
    BST<Integer> bst = new BST<>();
    int[] nums = {5, 3, 6, 8, 4, 2};
    for (int num : nums) {
      bst.add(num);
    }
  }

创建之后的二叉树

二叉树遍历

二叉树遍历方式，以上图为例

• 深度遍历
  前序遍历，根节点->左孩子->右孩子 5->3->2->4->6->8
  中序遍历，左孩子->根节点->右孩子 2->3->4->5->6->8
  后序遍历，左孩子->右孩子->根节点 2->4->3->8->6->5
• 广度遍历
  层序遍历，按层级遍历 5->3->6->2->4->8

前序遍历（递归）

  /**
   * 递归实现前序遍历
   */
  public List<E> preOrder() {
    List<E> res = new ArrayList<>();
    preOrder(root, res);
    return res;
  }

  private void preOrder(Node<E> root, List<E> res) {
    if (root != null) {
      res.add(root.data);
      preOrder(root.left, res);
      preOrder(root.right, res);
    }
  }

中序遍历（递归）

  /**
   * 递归实现中序遍历
   */
  public List<E> inOrder() {
    List<E> res = new ArrayList<>();
    inOrder(root, res);
    return res;
  }

  private void inOrder(Node<E> root, List<E> res) {
    if (root != null) {
      inOrder(root.left, res);
      res.add(root.data);
      inOrder(root.right, res);
    }
  }

后续遍历（递归）

/**
   * 递归实现后序遍历
   */
  public List<E> postOrder() {
    List<E> res = new ArrayList<>();
    postOrder(root, res);
    return res;
  }

  private void postOrder(Node<E> root, List<E> res) {
    if (root != null) {
      postOrder(root.left, res);
      postOrder(root.right, res);
      res.add(root.data);
    }
  }

可以看到使用递归的方式，3中遍历的写法基本一致，非递归写法相比递归要难一些，所以面试一般会问非递归写法。

前序遍历（非递归）

  /**
   * 非递归实现前序遍历
   */
  public List<E> preOrderNonRecursive() {
    List<E> res = new ArrayList<>();
    Stack<Node<E>> stack = new Stack<>();
    //将根节点入栈
    stack.push(root);
    while (!stack.isEmpty()) {
      Node<E> cur = stack.pop();
      if (cur != null) {
        //访问当前节点
        res.add(cur.data);
        //将右孩子入栈
        stack.push(cur.right);
        //将左孩子入栈
        stack.push(cur.left);
      }
    }
    return res;
  }

以上图为例，前序遍历过程为

• 5入栈
• 5出栈，访问5节点，右孩子6入栈，左孩子3入栈
• 3出栈，访问，右孩子4入栈，左孩子2入栈
• 2出栈，访问，右孩子空入栈，左孩子空入栈
• 出栈，栈顶为空，继续出栈
• 4出栈，访问，右孩子空入栈，左孩子空入栈
• 出栈，栈顶为空，继续出栈
• 6出栈，访问，右孩子8入栈，左孩子空入栈
• 出栈，栈顶为空，继续出栈
• 8出栈，访问，右孩子空入栈，左孩子空入栈
• 出栈，栈顶为空，继续出栈，栈为空，遍历结束

中序遍历（非递归）

  /**
   * 非递归实现中序遍历
   */
  public List<E> inOrderNonRecursive() {
    List<E> res = new ArrayList<>();
    Stack<Node<E>> stack = new Stack<>();
    Node<E> cur = root;
    while (cur != null || !stack.isEmpty()) {
      //将最左节点入栈
      while (cur != null) {
        stack.push(cur);
        cur = cur.left;
      }
      cur = stack.pop();
      res.add(cur.data);
      cur = cur.right;
    }
    return res;
  }

后序遍历（非递归）

  /** 
   * 非递归实现后序遍历 经典写法
   */
  public List<E> postOrderNonRecursive2() {
    List<E> res = new ArrayList<>();
    Stack<Node<E>> stack = new Stack<>();
    //当前访问节点
    Node<E> cur = root;
    //上一次访问节点
    Node<E> pre = null;
    while (cur != null || !stack.isEmpty()) {
      //将最左节点入栈
      while (cur != null) {
        stack.push(cur);
        cur = cur.left;
      }
      cur = stack.peek();
      //防止右孩子节点重复遍历
      if (cur.right != null && cur.right != pre) {
        cur = cur.right;
      } else {
        cur = stack.pop();
        res.add(cur.data);
        pre = cur;
        //置空 防止左孩子节点重复入栈
        cur = null;
      }
    }
    return res;
  }

后序遍历还是很麻烦的，这里有一种取巧的方式，前序遍历为根左右，后序为左右根，那么我们按根右左遍历，然后反转就是左右根。

  /**
   * 非递归实现后序遍历（取巧方式，根右左遍历，取反）
   */
  public List<E> postOrderNonRecursive() {
    List<E> res = new ArrayList<>();
    Stack<Node<E>> stack = new Stack<>();
    //将根节点入栈
    stack.push(root);
    while (!stack.isEmpty()) {
      Node<E> cur = stack.pop();
      if (cur != null) {
        //访问当前节点
        res.add(cur.data);
        //将左孩子入栈
        stack.push(cur.left);
        //将右孩子入栈
        stack.push(cur.right);
      }
    }
    //列表反转
    Collections.reverse(res);
    return res;
  }

模拟系统栈实现前中后序遍历

非递归的写法相比递归写法更加的困难，写法也不统一，那么有没有统一的写法呢，也是有的，那就是模拟系统栈。

  private static class Command<E> {

    static final String TYPE_GO = "go";
    static final String TYPE_PRINT = "print";
    /**
     * 表示指令，主要有两种：go print
     */
    String type;
    /**
     * 表示指令作用的节点
     */
    Node<E> node;

    Command(String type, Node<E> node) {
      this.type = type;
      this.node = node;
    }
  }

/**
   * 通过模拟系统栈的方式实现前序遍历
   */
  public List<E> preOrderMockSystem() {
    List<E> res = new ArrayList<>();
    Stack<Command<E>> stack = new Stack<>();
    stack.push(new Command<>(Command.TYPE_GO, root));//表示的指令是去根节点
    while (!stack.isEmpty()) {
      Command<E> command = stack.pop();
      if (command.node == null) {
        continue;
      }
      if (Command.TYPE_PRINT.equals(command.type)) {
        res.add(command.node.data);
      } else {// 指令是go
        stack.push(new Command<>(Command.TYPE_GO, command.node.right));
        stack.push(new Command<>(Command.TYPE_GO, command.node.left));
        stack.push(new Command<>(Command.TYPE_PRINT, command.node));
      }
    }
    return res;
  }

  /**
   * 通过模拟系统栈的方式实现中序遍历
   */
  public List<E> inOrderMockSystem() {
    List<E> res = new ArrayList<>();
    Stack<Command<E>> stack = new Stack<>();
    stack.push(new Command<>(Command.TYPE_GO, root));//表示的指令是去根节点
    while (!stack.isEmpty()) {
      Command<E> command = stack.pop();
      if (command.node == null) {
        continue;
      }
      if (Command.TYPE_PRINT.equals(command.type)) {
        res.add(command.node.data);
      } else {// 指令是go
        stack.push(new Command<>(Command.TYPE_GO, command.node.right));
        stack.push(new Command<>(Command.TYPE_PRINT, command.node));
        stack.push(new Command<>(Command.TYPE_GO, command.node.left));
      }
    }
    return res;
  }

  /**
   * 通过模拟系统栈的方式实现后序遍历
   */
  public List<E> postOrderMockSystem() {
    List<E> res = new ArrayList<>();
    Stack<Command<E>> stack = new Stack<>();
    stack.push(new Command<>(Command.TYPE_GO, root));//表示的指令是去根节点
    while (!stack.isEmpty()) {
      Command<E> command = stack.pop();
      if (command.node == null) {
        continue;
      }
      if (Command.TYPE_PRINT.equals(command.type)) {
        res.add(command.node.data);
      } else {// 指令是go
        stack.push(new Command<>(Command.TYPE_PRINT, command.node));
        stack.push(new Command<>(Command.TYPE_GO, command.node.right));
        stack.push(new Command<>(Command.TYPE_GO, command.node.left));
      }
    }
    return res;
  }

层序遍历

  /**
   * 层序遍历
   */
  public List<E> levelOrder() {
    List<E> res = new ArrayList<>();
    Queue<Node<E>> queue = new LinkedList<>();
    //将根节点入队
    queue.add(root);
    while (!queue.isEmpty()) {
      Node<E> cur = queue.poll();
      if (cur != null) {
        //访问当前节点
        res.add(cur.data);
        //将左孩子入队
        queue.add(cur.left);
        //将右孩子入队
        queue.add(cur.right);
      }
    }
    return res;
  }

删除任意节点

删除节点一共有4种情况

• 待删除节点有左孩子和右孩子
• 待删除节点有左孩子，无右孩子
• 待删除节点无左孩子，有右孩子
• 待删除节点无左孩子，无右孩子

第一种情况我们也可以通过后继节点转换成后面3种情况。

还是以上图为例，删除5节点的过程为

• 查找到5节点
• 5节点有左孩子和右孩子，查找5节点的后继节点6节点，使用6节点数据替换5节点数据，接下来删除6节点就可以了
• 将6节点的左孩子或右孩子和6节点的父节点连接上，这样6节点就删除了

  /**
   * 删除指定元素
   */
  public boolean remove(E e) {
    Node<E> node = find(root, e);
    if (node == null) {
      return false;
    }
    fastRemove(node);
    size--;
    return true;
  }

  /**
   * 在以root为根节点的树中查找值为e的节点
   */
  private Node<E> find(Node<E> root, E e) {
    if (root != null) {
      if (e.compareTo(root.data) > 0) {
        return find(root.right, e);
      } else if (e.compareTo(root.data) < 0) {
        return find(root.left, e);
      } else {
        return root;
      }
    }
    return root;
  }

  private void fastRemove(Node<E> node) {
    //node为待删除的节点
    //将删除一个有左孩子和右孩子的节点的情况转换成删除没有左孩子或没有右孩子的情况
    //查找待删除节点的后继节点
    if (node.left != null && node.right != null) {
      //使用后继节点代替待删除节点
      Node<E> successor = minimum(node.right);
      node.data = successor.data;
      node = successor;
    }
    Node<E> replacement = (node.left != null) ? node.left : node.right;
    if (replacement != null) {
      replacement.parent = node.parent;
    }
    if (node.parent == null) {
      //待删除节点没有父节点
      root = replacement;
    } else {
      if (node == node.parent.left) {
        node.parent.left = replacement;
      } else {
        node.parent.right = replacement;
      }
    }
    node.left = node.right = node.parent = null;
  }

  private Node<E> minimum(Node<E> root) {
    Node<E> cur = root;
    while (cur.left != null) {
      cur = cur.left;
    }
    return cur;
  }

上面的代码实现是依赖于节点中有指向父节点的指针，如果没有父节点指针的话，就要换一种写法了。

  /**
   * 删除任意元素
   */
  public void removeWithoutParent(E e) {
    boolean success = fastRemoveWithoutParent(root, e);
    if (success) {
      size--;
    }
  }

  /**
   * 删除以node为根节点中值为e的节点
   *
   * @param node 根节点
   * @param e 待删除的值
   * @return 删除是否成功
   */
  private boolean fastRemoveWithoutParent(Node<E> node, E e) {
    //待删除节点的父节点
    Node<E> parent = null;
    Node<E> cur = node;
    //查找到待删除的节点
    while (cur != null) {
      if (e.compareTo(cur.data) > 0) {
        parent = cur;
        cur = cur.right;
      } else if (e.compareTo(cur.data) < 0) {
        parent = cur;
        cur = cur.left;
      } else {
        break;
      }
    }
    //表示没有找到待删除节点
    if (cur == null) {
      return false;
    }
    //将删除一个有左孩子和右孩子的节点的情况转换成删除没有左孩子或没有右孩子的情况
    //查找待删除节点的后继节点
    if (cur.left != null && cur.right != null) {
      Node<E> successor = cur.right;
      parent = cur;
      while (successor.left != null) {
        parent = successor;
        successor = successor.left;
      }
      //使用后继节点代替待删除节点
      cur.data = successor.data;
      cur = successor;
    }

    Node<E> replacement = (cur.left != null) ? cur.left : cur.right;
    if (parent == null) {
      //待删除节点没有父节点
      root = replacement;
    } else {
      if (cur == parent.left) {
        parent.left = replacement;
      } else {
        parent.right = replacement;
      }
    }
    cur.left = cur.right = null;
    return true;
  }

在遍历和查找后继节点的过程中，查找到待删除节点的父节点，之后的逻辑就和上面的一样了。

是否包含指定元素

  /**
   * 是否包含指定元素
   */
  public boolean contains(E e) {
    Node<E> node = find(root, e);
    return node != null;
  }

总结

二叉搜索树支持快速的查找、插入、删除操作，jdk中的TreeMap就是一种更优的二叉搜索树实现（平衡二叉树）。关于代码中使用到的栈和队列，不了解的可以看 java使用数组和链表实现栈和队列