Count of Smaller Number before itself

Give you an integer array (index from 0 to n-1, where n is the size of this array, value from 0 to 10000) . For each element Ai in the array, count the number of element before this element Ai is smaller than it and return count number array.

Example:

For array [1,2,7,8,5], return [0,1,2,3,2]

Solution:

This is a variation of Interval Sum II. We could create a segment tree with range from 0 to 10000. For each range (0,0 or 1,1) in segment tree, we have count 0. When a new number A[i] comes in, we modify its count in the segment tree and query the interval sum from 0 to A[i] - 1 as the i-th output.

Code:

public class Solution {
   /**
     * @param A: An integer array
     * @return: Count the number of element before this element 'ai' is 
     *          smaller than it and return count number array
     */ 

    public ArrayList<Integer> countOfSmallerNumberII(int[] A) {
        ArrayList<Integer> res = new ArrayList<>();
        if (A == null || A.length == 0) {
            return res;
        }
        SegmentNode root = build(0, 10000);
        for (int i = 0; i < A.length; i++) {
            int count = 0;
            if (A[i] > 0) {
                count = query(root, 0, A[i] - 1);
            }
            modify(root, A[i], 1);
            res.add(count);
        }
        return res;
    }
    private SegmentNode build(int start, int end) {
        Stack<SegmentNode> stack = new Stack<>();
        SegmentNode root = new SegmentNode(start, end);
        stack.push(root);
        while (!stack.empty()) {
            SegmentNode node = stack.pop();
            if (node.start >= node.end) {
                continue;
            }
            int mid = node.start + (node.end - node.start) / 2;

            if (node.start <= mid) {
                node.left = new SegmentNode(node.start, mid);
                stack.push(node.left);
            }
            if (mid + 1 <= node.end) {
                node.right = new SegmentNode(mid + 1, node.end);
                stack.push(node.right);
            }
        }
        return root;
    }

    private int query(SegmentNode node, int start, int end) {
        if (node == null || start > end) {
            return 0;
        }
        if (start == node.start && node.end == end) {
            return node.count;
        }
        int total = 0;
        int mid = node.start + (node.end - node.start) / 2;
        if (start <= mid) {
            if (mid < end) {
                total += query(node.left, start, mid);
            } else {
                total += query(node.left, start, end);
            }
        } 
        if (mid < end) {
            if (start <= mid) {
                total += query(node.right, mid + 1, end);
            } else {
                total += query(node.right, start, end);
            }
        }
        return total;
    }
    private void modify(SegmentNode node, int index, int count) {
        if (node == null) {
            return;
        }
        if (node.start == index && node.end == index) {
            node.count += count;
            return;
        }
        int mid = node.start + (node.end - node.start) / 2;
        if (index <= mid) {
            modify(node.left, index, count);
        } else {
            modify(node.right, index, count);
        }
        if (node.left != null) {
            node.count = node.left.count;
        }
        if (node.right != null) {
            node.count += node.right.count;
        }
    }

    class SegmentNode {
        int start;
        int end;
        int count;
        SegmentNode left;
        SegmentNode right;
        public SegmentNode(int start, int end) {
            this.start = start;
            this.end = end;
            this.count = 0;
        }
    }
}