Count of Smaller Numbers After Self [Hard] — Merge Sort

Problem Statement

Given integer array nums, return a counts array where counts[i] is the number of smaller elements to the right of nums[i].

Input: [5,2,6,1] → Output: [2,1,1,0]

Intuition

Modified merge sort: during merge, when we pick an element from the right half, all remaining left-half elements have a smaller element to the right (count their right-half contributions).

Solutions

Python

def countSmaller(nums: list[int]) -> list[int]:
    n = len(nums)
    counts = [0] * n
    indexed = list(enumerate(nums))

    def merge_sort(arr):
        if len(arr) <= 1:
            return arr
        mid = len(arr) // 2
        left  = merge_sort(arr[:mid])
        right = merge_sort(arr[mid:])
        merged = []
        i = j = 0
        while i < len(left) and j < len(right):
            if left[i][1] > right[j][1]:
                # left[i] is greater than right[j]
                # all remaining right elements until j that were already merged count
                counts[left[i][0]] += len(right) - j
                merged.append(left[i]); i += 1
            else:
                merged.append(right[j]); j += 1
        while i < len(left):
            counts[left[i][0]] += len(right) - j
            merged.append(left[i]); i += 1
        while j < len(right):
            merged.append(right[j]); j += 1
        return merged

    merge_sort(indexed)
    return counts

C++

vector<int> counts;
void mergeSort(vector<pair<int,int>>& arr, int l, int r) {
    if (l>=r) return;
    int mid=(l+r)/2;
    mergeSort(arr,l,mid); mergeSort(arr,mid+1,r);
    vector<pair<int,int>> tmp;
    int i=l,j=mid+1;
    while (i<=mid&&j<=r) {
        if (arr[i].first>arr[j].first) {
            counts[arr[i].second]+=r-j+1;
            tmp.push_back(arr[i++]);
        } else tmp.push_back(arr[j++]);
    }
    while(i<=mid) tmp.push_back(arr[i++]);
    while(j<=r)   tmp.push_back(arr[j++]);
    for(int k=l;k<=r;k++) arr[k]=tmp[k-l];
}
vector<int> countSmaller(vector<int>& nums) {
    int n=nums.size();
    counts.assign(n,0);
    vector<pair<int,int>> arr;
    for(int i=0;i<n;i++) arr.push_back({nums[i],i});
    mergeSort(arr,0,n-1);
    return counts;
}

Complexity

Time: O(n log n)
Space: O(n)

Alternative: BIT/Fenwick Tree

Coordinate compress, then process right to left querying BIT for count of smaller elements.