Share My AVL Tree Solution, O(NlgN) time.


  • 1
    H

    Ref:
    (1) http://www.geeksforgeeks.org/avl-tree-set-1-insertion/

    public class Solution {

    public int reversePairs(int[] nums) {
        
        // Algo thinking: building a BST, go left when node.val <= 2 * root.val, right otherwise
        // But need to keep it balanced -> AVL Tree or Red-Black Tree
        // time = O(NlgN), space = O(N)
        
        if (nums == null || nums.length == 0) return 0;
        
        int n = nums.length;
        
        TreeNode root = new TreeNode(nums[0]);
        int ans = 0;
        for (int i = 1; i < nums.length; i++) {
            ans += search(root, (long) nums[i] * 2);
            root = insert(root, (long) nums[i]);
    
            // preOrder(root);
            // System.out.println();
        }
        
        return ans;
            
    }
    
    private int search(TreeNode root, long key) {
        
        if (root == null) return 0;
        
        if (key < root.val) {       // key < root.val:  go left
            return root.rightCount + search(root.left, key);
        } else {                    // key >= root.val: go right
            return search(root.right, key);
        }
    }
    
    private TreeNode insert(TreeNode root, long key) {
        
        if (root == null) return new TreeNode(key);
        
        if (key < root.val) {   // key < root.val:  go left
            root.left = insert(root.left, key);
        } else if (key == root.val){
            root.rightCount++;
            return root;
        } else {
            root.rightCount++;
            root.right = insert(root.right, key);
        }
        
        root.height = Math.max(getHeight(root.left), getHeight(root.right)) + 1;
        
        int balance = getBalance(root);
        
        // System.out.println(root.val + " balance " + balance);
        
        // case 1 left left 
        if (balance > 1 && getHeight(root.left.left) > getHeight(root.left.right)) {
            return rightRotate(root);
        }
        
        // case 2 left right 
        if (balance > 1 && getHeight(root.left.left) < getHeight(root.left.right)) {
            root.left = leftRotate(root.left);
            return  rightRotate(root);
        }
        
        // case 3 right right
        if (balance < -1 && getHeight(root.right.left) < getHeight(root.right.right)) {
            return leftRotate(root);
        }
        
        // case 4 right left 
        if (balance < -1 && getHeight(root.right.left) > getHeight(root.right.right)) {
            root.right = rightRotate(root.right);
            return leftRotate(root);
        }
        
        return root;
    }
    
    private TreeNode leftRotate(TreeNode root) {
        
        // setp 1: take care of nodes
        TreeNode newRoot = root.right;
        TreeNode b = newRoot.left;
        
        newRoot.left = root;
        root.right = b;
        
        // step 2: take care of height
        root.height = Math.max(getHeight(root.left), getHeight(root.right)) + 1;
        newRoot.height = Math.max(getHeight(newRoot.left), getHeight(newRoot.right)) + 1;
        
        // step 3: take care of rightCount
        root.rightCount -= getRightCount(newRoot);
    
        return newRoot;
    }
    
    private TreeNode rightRotate(TreeNode root) {
        
        // setp 1: take care of nodes
        TreeNode newRoot = root.left;
        TreeNode b = newRoot.right;
        
        newRoot.right = root;
        root.left = b;
        
        // step 2: take care of height
        root.height = Math.max(getHeight(root.left), getHeight(root.right)) + 1;
        newRoot.height = Math.max(getHeight(newRoot.left), getHeight(newRoot.right)) + 1;
        
        // step 3: take care of rightCount
        newRoot.rightCount += getRightCount(root);
        
        return newRoot;
    }
    
    
    private int getHeight(TreeNode node) {
        return node == null ? 0 : node.height;
    }
    
    private int getBalance(TreeNode node) {
        return node == null ? 0 : getHeight(node.left) - getHeight(node.right);
    }
    
    private int getRightCount(TreeNode node) {
        return node == null ? 0 : node.rightCount;
    }
    
    private void preOrder(TreeNode root) {
        
        if (root == null) {
            System.out.print("NIL ");
            return;
        }
        
        System.out.print(root.val + " ");
        preOrder(root.left);
        preOrder(root.right);
    }
    
    class TreeNode {
        
        long val;
        int rightCount;
        int height;
        TreeNode left;
        TreeNode right;
        public TreeNode(long val) {
            this.val = val;
            height = 1;
            rightCount = 1;
        }
    }
    

    }


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.