Elegant BST solution using order statistics


  • 0
    V
    class TreeNode {
        TreeNode left;
        TreeNode right;
        int value;
        int leftCount; // number of nodes in the left subtree
        int duplicate; // how many times value has been duplicated, default to 1
         
        public TreeNode(int value) {
            this.value = value;
            this.leftCount = 0;
            this.duplicate = 1;
        }
    }
     
    public class Solution {
        public List<Integer> countSmaller(int[] nums) {
            return bst_Solution(nums);
        }
         
        private List<Integer> bst_Solution(int[] nums) {
            int[] inversions = new int[nums.length];
            TreeNode root = null;
            for(int i = nums.length - 1; i >= 0; i--) {
                root = bstInsert(root, i, nums, inversions);
            }
             
            List<Integer> res = new ArrayList<>();
            for(int i = 0; i < inversions.length; i++) {
                res.add(inversions[i]);
            }
             
            return res;
        }
         
        // elegant BST insert with order-statistics and inversion counting
        private TreeNode bstInsert(TreeNode root, int idx, int[] nums, int[] inversions) {
            if (root == null) {
                return new TreeNode(nums[idx]);
            }
             
            if (root.value == nums[idx]) {
                inversions[idx] += root.leftCount;
                root.duplicate++;
            } else if (nums[idx] < root.value) {
                root.leftCount++;
                root.left = bstInsert(root.left, idx, nums, inversions);
            } else {
                inversions[idx] += root.leftCount + root.duplicate;
                root.right = bstInsert(root.right, idx, nums, inversions);
            }
             
            return root;
        }
    }

Log in to reply
 

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