LeetCode - Count of Smaller Numbers After Self

CountofSmallerNumbersAfterSelf.java

// You are given an integer array nums and you have to return a new counts array. The counts array has the property where counts[i] is the number of smaller elements to the right of nums[i].

// Example:

// Given nums = [5, 2, 6, 1]

// To the right of 5 there are 2 smaller elements (2 and 1).
// To the right of 2 there is only 1 smaller element (1).
// To the right of 6 there is 1 smaller element (1).
// To the right of 1 there is 0 smaller element.
// Return the array [2, 1, 1, 0].

public class Solution {
    
    class SegmentTreeNode{
        int start, end, sum;
        SegmentTreeNode left, right;
        SegmentTreeNode(int s, int e, int sum){
            this.start = s;
            this.end = e;
            this.sum = sum;
        }
    }
    
    private SegmentTreeNode root;
    
    private SegmentTreeNode buildTree(int start, int end){
        if(start > end) return null;
        if(start == end) return new SegmentTreeNode(start, end, 0);
        int mid = start + (end - start) / 2;
        SegmentTreeNode left = buildTree(start, mid);
        SegmentTreeNode right = buildTree(mid + 1, end);
        SegmentTreeNode root = new SegmentTreeNode(start, end, left.sum + right.sum);
        root.left = left; root.right = right;
        return root;
    }
    
    private void modify(SegmentTreeNode root, int index){
        if(root == null) return;
        if(root.start == index && root.end == index){
            root.sum += 1;
            return;
        }
        
        int mid = root.start + (root.end - root.start) / 2;
        if(root.start <= index && index <= mid){
            modify(root.left, index);
        } else if(root.end >= index && index > mid){
            modify(root.right, index);
        }
        root.sum = root.left.sum + root.right.sum;
    }
    
    private int query(SegmentTreeNode root, int start, int end){
        if(root == null) return 0;
        if(start <= root.start && end >= root.end) return root.sum;
        int mid = root.start + (root.end - root.start) / 2;
        if(end <= mid){
            return query(root.left, start, end);
        } else if(start > mid){
            return query(root.right, start, end);
        } else{
            return query(root.left, start, end) + query(root.right, start, end);
        }
    }
    
    public List<Integer> countSmaller(int[] nums) {
        List<Integer> res = new ArrayList<Integer>();
        if(nums == null || nums.length == 0) return res;
        int start = Integer.MAX_VALUE, end = Integer.MIN_VALUE;
        for(int num : nums){
            start = Math.min(num, start);
            end = Math.max(num, end);
        }
        root = buildTree(start, end);
        
        for(int i = nums.length - 1; i >= 0; i--){
            res.add(0, query(root, start, nums[i] - 1));
            modify(root, nums[i]);
        }
        
        return res;
    }
}

// solution 2 --- binary search tree
public List<Integer> countSmaller(int[] nums) {
    List<Integer> res = new ArrayList<Integer>();
    if(nums == null || nums.length == 0) return res;
    TreeNode root = new TreeNode(nums[nums.length - 1]);
    res.add(0);
    for(int i = nums.length - 2; i >= 0; i--){
        res.add(0, insert(root, nums[i]));
    }
    
    return res;
}

private int insert(TreeNode root, int value){
    int res = 0;
    while(true){
        if(root.val == value){
            res += root.leftCount;
            root.selfCount++;
            break;
        } else if(root.val > value){
            root.leftCount++;
            if(root.left == null){
                root.left = new TreeNode(value);
                break;
            } else{
                root = root.left;
            }
        } else{
            res += root.selfCount + root.leftCount;
            if(root.right == null){
                root.right = new TreeNode(value);
                break;
            } else{
                root = root.right;
            }
        }
    }
    return res;
}

class TreeNode{
    int val, leftCount, selfCount;
    TreeNode left, right;
    TreeNode(int value){
        this.val = value;
        selfCount = 1;
    }
}

// solution 3 --- merge sort 

// solution 4 --- binary indexed tree