平衡二叉树java实现

写完之后，整个人心情都不太好

单元测试

package com.fly.study.structure;

import org.junit.Before;
import org.junit.Test;
import org.omg.PortableInterceptor.INACTIVE;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.IntStream;

import static org.junit.Assert.*;

/**
 * @Auther: fly
 * @Date: 2018/10/21 17:16
 * @see: https://github.com/FightingIsLife/structure
 * @version: 1
 * @email: [email protected]
 * @Description:
 */
public class BinaryTreeTest {

    private BinaryTree binaryTree = new BinaryTree<>(Integer::compare);

    @Test
    public void toListByPreOrder() {
        add();
        System.out.println(binaryTree.toListByPreOrder());
    }

    @Test
    public void toListByInOrder() {
        add();
        System.out.println(binaryTree.toListByInOrder());
    }

    @Test
    public void toListByPostOrder() {
        add();
        System.out.println(binaryTree.toListByPostOrder());
    }

    @Test
    public void add() {
        List nums = Arrays.asList(4, 1, 3, 10, 7, 8, 9, 5, 2, 6);
        for (int num : nums) {
            binaryTree.add(num);
        }
        Collections.sort(nums);
        assertEquals(nums, binaryTree.toListByInOrder());
    }

    @Test
    public void remove() {
        add();
        assertTrue(binaryTree.remove(7));
        assertTrue(binaryTree.remove(9));
        assertTrue(binaryTree.remove(1));
        assertEquals(Arrays.asList(2, 3, 4, 5, 6, 8, 10), binaryTree.toListByInOrder());
    }

    @Test
    public void contains() {
        IntStream.range(1, 100).filter(this::isEven).forEach(binaryTree::add);
        assertTrue(binaryTree.contains(2));
        assertTrue(binaryTree.contains(98));
        assertFalse(binaryTree.contains(0));
        assertFalse(binaryTree.contains(99));
    }

    private boolean isEven(int num) {
        return num % 2 == 0;
    }

    @Test
    public void getSize() {
        assertEquals(0, binaryTree.getSize());
        binaryTree.add(1);
        assertEquals(1, binaryTree.getSize());
        binaryTree.add(2);
        assertEquals(2, binaryTree.getSize());
        binaryTree.add(2);
        assertEquals(2, binaryTree.getSize());
        binaryTree.remove(2);
        assertEquals(1, binaryTree.getSize());
        binaryTree.remove(0);
        assertEquals(1, binaryTree.getSize());
    }
}

实现

package com.fly.study.structure;

import java.util.*;

import static java.lang.Integer.max;

/**
 * @Auther: fly
 * @Date: 2018/10/21 12:58
 * @see: https://github.com/FightingIsLife/structure
 * @version: 1
 * @email: [email protected]
 * @Description:
 */
public class BinaryTree {

    private Comparator comparator;
    private int size = 0;
    private Node root;

    /**
     * 创建二叉树
     *
     * @param comparator 比较器
     */
    public BinaryTree(Comparator comparator) {
        this.comparator = comparator;
    }

    /**
     * 获取先序遍历集合
     *
     * @return
     */
    public List toListByPreOrder() {
        List list = new ArrayList<>(size);
        toListByPreOrder(root, list);
        return list;
    }

    private void toListByPreOrder(Node node, List list) {
        if (node != null) {
            list.add(node.data);
            toListByPreOrder(node.leftNode, list);
            toListByPreOrder(node.rightNode, list);
        }
    }

    /**
     * 获取中序遍历集合
     *
     * @return
     */
    public List toListByInOrder() {
        List list = new ArrayList<>(size);
        toListByInOrder(root, list);
        return list;
    }

    private void toListByInOrder(Node node, List list) {
        if (node != null) {
            toListByInOrder(node.leftNode, list);
            list.add(node.data);
            toListByInOrder(node.rightNode, list);
        }
    }

    /**
     * 获取后序遍历集合
     *
     * @return
     */
    public List toListByPostOrder() {
        List list = new ArrayList<>(size);
        toListByPostOrder(root, list);
        return list;
    }

    private void toListByPostOrder(Node node, List list) {
        if (node != null) {
            toListByInOrder(node.leftNode, list);
            toListByInOrder(node.rightNode, list);
            list.add(node.data);
        }
    }

    /**
     * 向二叉树添加一个元素
     *
     * @param data
     * @return true添加成功，false添加失败
     */
    public boolean add(T data) {
        boolean success = true;
        if (root == null) {
            root = new Node<>();
            root.data = data;
        } else {
            success = add(root, data);
        }
        if (success) {
            size++;
        }
        return success;
    }

    private boolean add(Node node, T data) {
        int compare = comparator.compare(data, node.data);
        boolean success = true;
        if (compare < 0) {
            // 往左子树加
            if (node.leftNode == null) {
                setLeftNode(node, new Node<>(data));
            } else {
                success = add(node.leftNode, data);
            }
        } else if (compare > 0) {
            // 往右子树加
            if (node.rightNode == null) {
                setRightNode(node, new Node<>(data));
            } else {
                success = add(node.rightNode, data);
            }
        } else {
            // compare相等的值，不再重复添加
            success = false;
        }
        updateAndBalanceNode(node);
        return success;
    }

    /**
     * 向二叉树中删除一个元素
     *
     * @param data 元素
     * @return
     */
    public boolean remove(T data) {
        boolean success = removeByData(data);
        if (success) {
            size--;
        }
        return success;
    }

    public boolean removeByData(T data) {
        Node node = search(root, data);
        if (node == null) {
            return false;
        }
        Node balanceNode = node.parentNode;
        if (node.deep == 0) {
            setParentNode(null, node.parentNode, isLeftNode(node.parentNode, node));
            if (root == node) {
                root = null;
            }
        } else if (node.deepDiff() > 0) {
            Node rightMostNode = rightMost(node.leftNode);
            balanceNode = rightMostNode;
            if (node.leftNode != rightMostNode) {
                balanceNode = rightMostNode.parentNode;
                setLeftNode(rightMostNode, node.leftNode);
            }
            setRightNode(rightMostNode, node.rightNode);
            setParentNode(null, rightMostNode.parentNode, isLeftNode(rightMostNode, rightMostNode.parentNode));
            setParentNode(rightMostNode, node.parentNode, isLeftNode(node.parentNode, node));
            if (root == node) {
                root = rightMostNode;
            }
        } else {
            Node leftMostNode = leftMost(node.rightNode);
            balanceNode = leftMostNode;
            if (node.rightNode != leftMostNode) {
                balanceNode = leftMostNode.parentNode;
                setRightNode(leftMostNode, node.rightNode);
            }
            setLeftNode(leftMostNode, node.leftNode);
            setParentNode(null, leftMostNode.parentNode, isLeftNode(leftMostNode, leftMostNode.parentNode));
            setParentNode(leftMostNode, node.parentNode, isLeftNode(node.parentNode, node));
            if (root == node) {
                root = leftMostNode;
            }
        }
        reBalance(balanceNode);
        return true;
    }

    private Node rightMost(Node node) {
        if (node == null) {
            return null;
        }
        while (node.rightNode != null) {
            node = node.rightNode;
        }
        return node;
    }

    private Node leftMost(Node node) {
        if (node == null) {
            return null;
        }
        while (node.leftNode != null) {
            node = node.leftNode;
        }
        return node;
    }

    private void reBalance(Node node) {
        if (node != null) {
            Node parentNode = node.parentNode;
            updateAndBalanceNode(node);
            reBalance(parentNode);
        }
    }

    /**
     * 查询二叉树中是否包含data
     */
    public boolean contains(T data) {
        return search(root, data) != null;
    }

    private Node search(Node node, T data) {
        if (node != null) {
            int compare = comparator.compare(data, node.data);
            if (compare == 0) {
                return node;
            } else if (compare < 0) {
                return search(node.leftNode, data);
            } else {
                return search(node.rightNode, data);
            }
        }
        return null;
    }


    private void updateAndBalanceNode(Node node) {
        node.updateDeep();
        balance(node);
    }

    /**
     * 使二叉树的node节点平衡
     *
     * @param node 要平衡的节点
     */
    private void balance(Node node) {
        int deepDiff = node.deepDiff();
        if (deepDiff == 2) {
            if (node.leftNode.deepDiff() == -1) {
                leftRotate(node.leftNode);
            }
            rightRotate(node);
        } else if (deepDiff == -2) {
            if (node.rightNode.deepDiff() == 1) {
                rightRotate(node.rightNode);
            }
            leftRotate(node);
        }
    }

    private void leftRotate(Node node) {
        Node rightNode = node.rightNode;
        Node parentNode = node.parentNode;
        setRightNode(node, rightNode.leftNode);
        setLeftNode(rightNode, node);
        setParentNode(rightNode, parentNode, isLeftNode(parentNode, node));
        if (node == root) {
            root = rightNode;
        }
    }

    private void rightRotate(Node node) {
        Node leftNode = node.leftNode;
        Node parentNode = node.parentNode;
        setLeftNode(node, leftNode.rightNode);
        setRightNode(leftNode, node);
        setParentNode(leftNode, parentNode, isLeftNode(parentNode, node));
        if (node == root) {
            root = leftNode;
        }
    }

    private boolean isLeftNode(Node parentNode, Node node) {
        return parentNode != null && parentNode.leftNode == node;
    }

    private void setLeftNode(Node node, Node leftNode) {
        if (leftNode != null) {
            leftNode.parentNode = node;
        }
        if (node != null) {
            node.leftNode = leftNode;
            node.updateDeep();
        }
    }

    private void setRightNode(Node node, Node rightNode) {
        if (rightNode != null) {
            rightNode.parentNode = node;
        }
        if (node != null) {
            node.rightNode = rightNode;
            node.updateDeep();
        }
    }

    private void setParentNode(Node node, Node parentNode, boolean isLeft) {
        if (node != null) {
            node.parentNode = parentNode;
        }
        if (parentNode != null) {
            if (isLeft) {
                parentNode.leftNode = node;
            } else {
                parentNode.rightNode = node;
            }
        }
    }

    public static class Node {
        private Node parentNode;
        /**
         * 树的深度： deep = Max(leftNode.deep,rightNode.deep)+1
         *
         * @see this#updateDeep()
         */
        private int deep;
        private T data;
        private Node leftNode;
        private Node rightNode;

        public Node() {
        }

        public Node(T data) {
            this.data = data;
        }

        private void updateDeep() {
            deep = max(getDeep(leftNode), getDeep(rightNode)) + 1;
        }

        /**
         * 左右子树的深度差
         */
        private int deepDiff() {
            return getDeep(leftNode) - getDeep(rightNode);
        }

        private int getDeep(Node node) {
            return node == null ? -1 : node.deep;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "data=" + data +
                    '}';
        }
    }

    public int getSize() {
        return size;
    }

    @Override
    public String toString() {
        return "BinaryTree{" +
                "size=" + size +
                ", root=" + root +
                '}';
    }
}