【手撕代码】剑指offer-数据流的中位数

描述

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。我们使用Insert()方法读取数据流，使用GetMedian()方法获取当前读取数据的中位数。

解题思路

计算数据流的中位数难点：数据流是一个动态的数据，在求解中位数时需要考虑动态排序计算当前位置中位数。二叉平衡树可以针对数据流实现实现动态排序，但传统的二叉平衡树只会考虑左右子树的深度导致树失去平衡，本题所面对的场景则需要考虑左右字数节点数差值导致树失去平衡。因此插入新节点后我们需要检查是否因为插入新的节点从而导致树失去平衡，失去平衡的条件即为左右子树节点数差值大于1。树右旋时，需要找到左子树最大值并将其设为根节点。树左旋时，需要找到右子树最小值并将其设为根节点。

public class Solution {
    static class AvlNode {
        // 节点数
        int count;
        // 节点元素值
        int val;
        // 左孩子
        AvlNode left;
        // 右孩子
        AvlNode right;

        public AvlNode(int val) {
            this.val = val;
            this.count = 1;
        }
    }

    public static int getCount(AvlNode node) {
        return node == null ? 0 : node.count;
    }

    AvlNode root;

    public void Insert(Integer num) {
        root = Insert(root, num);
    }

    public AvlNode Insert(AvlNode node, int num) {
        if (node == null) {
            return new AvlNode(num);
        }
        if (num <= node.val) {
            node.left = Insert(node.left, num);
            node = balance(node, num);
        } else {
            node.right = Insert(node.right, num);
            node = balance(node, num);
        }
        node.count=getCount(node.left)+getCount(node.right)+1;
        return node;
    }

    public AvlNode balance(AvlNode node, int num) {
        // 左子树失去平衡
        if (getCount(node.left) - getCount(node.right) == 2) {
            if (num < node.left.val) {
                // LL型
                node = rotateWithLeftChild(node);
            } else {
                // LR型
                node = doubleRotateWithLeftChild(node);
            }
        }

        // 右子树失去平衡
        if (getCount(node.right) - getCount(node.left) == 2) {
            if (num > node.right.val) {
                // RR型
                node = rotateWithRightChild(node);
            } else {
                // RL型
                node = doubleRotateWithRightChild(node);
            }
        }

        return node;
    }

    public AvlNode rotateWithLeftChild(AvlNode node) {
        AvlNode leftChild = node.left;
        AvlNode parent = null;
        while (leftChild.right != null) {
            leftChild.count--;
            parent = leftChild;
            leftChild = leftChild.right;
        }
        if (leftChild != node.left) {
            // 保存最大值节点左孩子
            parent.right = leftChild.left;
            // 将最大值左孩子设置为根节点左孩子
            leftChild.left = node.left;
        }
        leftChild.right = node;
        node.left = null;
        // 更新右子树高度
        node.count=getCount(node.left)+getCount(node.right)+1;
        // 更新根节点高度
        leftChild.count=getCount(leftChild.left)+getCount(leftChild.right)+1;
        // 新的根节点
        return leftChild;
    }

    public AvlNode rotateWithRightChild(AvlNode node) {
        AvlNode rightChild = node.right;
        AvlNode parent = null;
        while (rightChild.left != null) {
            rightChild.count--;
            parent = rightChild;
            rightChild = rightChild.left;
        }
        if (rightChild != node.right) {
            // 保存最小值节点右孩子
            parent.left = rightChild.right;
            // 将最大值左孩子设置为根节点左孩子
            rightChild.right = node.right;
        }
        rightChild.left = node;
        node.right = null;
        // 更新左子树高度
        node.count=getCount(node.left)+getCount(node.right)+1;
        // 更新根节点高度
        rightChild.count=getCount(rightChild.left)+getCount(rightChild.right)+1;
        // 新的根节点
        return rightChild;
    }

    // 在左子树内部插入导致失去平衡-左右
    public AvlNode doubleRotateWithLeftChild(AvlNode node) {
        // 左子树左旋
        node.left = rotateWithRightChild(node.left);
        // 整棵树进行右旋
        node = rotateWithLeftChild(node);
        return node;
    }

    // 在右子树左侧插入导致失去平衡-右左
    public AvlNode doubleRotateWithRightChild(AvlNode node) {
        // 左子树右旋
        node.right = rotateWithLeftChild(node.right);
        // 整棵树进行左旋
        node = rotateWithRightChild(node);
        return node;
    }

    public Double GetMedian() {
        int leftCount =getCount(root.left);
        int rightCount = getCount(root.right);
        // 左右子树节点数相等,根节点为中位数
        if (leftCount == rightCount) {
            return root.val * 1.0;
        }
        // 左子树节点数多于右子树节点数
        if (leftCount > rightCount) {
            AvlNode node = root.left;
            while (node.right != null) {
                node = node.right;
            }
            return (node.val + root.val) / 2.0;
        } 
        // 右子树节点数多于左子树节点数
        else {
            AvlNode node = root.right;
            while (node != null && node.left != null) {
                node = node.left;
            }
            return (node.val + root.val) / 2.0;
        }
    }
}