Learning On the Way

315. Count of Smaller Numbers After Self

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].

Merge Sort

def countSmaller(self, nums: List[int]) -> List[int]:
    res = [0] * len(nums)
    
    array = list(zip(nums, range(len(nums))))
    
    def mergeSort(arr, left, right):
        if left == right:
            return
        
        mid = left + (right - left) // 2
        mergeSort(arr, left, mid)
        mergeSort(arr, mid + 1, right)
        merge(arr, left, mid, right)
        
    def merge(arr, left, mid, right):
        i = left
        j = mid + 1
        
        tmp = []
        
        while i <= mid or j <= right:
            if j > right or (i <= mid and arr[i][0] <= arr[j][0]):
                tmp.append(arr[i])
                res[arr[i][1]] += j - mid - 1
                i += 1
            else:
                tmp.append(arr[j])
                j += 1
                
        arr[left:right + 1] = tmp
        
    mergeSort(array, 0, len(nums) - 1)
    return res

Segment Tree

def countSmaller(self, nums: List[int]) -> List[int]: 
    # implement segment tree 
    def update(index, value, tree, size): 
        index += size  # shift the index to the leaf 
        # update from leaf to root 
        tree[index] += value 
        while index > 1: 
            index //= 2 
            tree[index] = tree[index * 2] + tree[index * 2 + 1] 
    def query(left, right, tree, size): 
        # return sum of [left, right) 
        result = 0 
        left += size  # shift the index to the leaf 
        right += size 
        while left < right: 
            # if left is a right node 
            # bring the value and move to parent's right node 
            if left % 2 == 1: 
                result += tree[left] 
                left += 1 
            # else directly move to parent 
            left //= 2 
            # if right is a right node 
            # bring the value of the left node and move to parent 
            if right % 2 == 1: 
                right -= 1 
                result += tree[right] 
            # else directly move to parent 
            right //= 2 
        return result 
    offset = 10**4  # offset negative to non-negative 
    size = 2 * 10**4 + 1  # total possible values in nums 
    tree = [0] * (2 * size) 
    result = [] 
    for num in reversed(nums): 
        smaller_count = query(0, num + offset, tree, size) 
        result.append(smaller_count) 
        update(num + offset, 1, tree, size) 
    return reversed(result)

Binary Index Tree

def countSmaller(self, nums: List[int]) -> List[int]: 
    # implement Binary Index Tree 
    def update(index, value, tree, size): 
        index += 1  # index in BIT is 1 more than the original index 
        while index < size: 
            tree[index] += value 
            index += index & -index 
    def query(index, tree): 
        # return sum of [0, index) 
        result = 0 
        while index >= 1: 
            result += tree[index] 
            index -= index & -index 
        return result 
    offset = 10**4  # offset negative to non-negative 
    size = 2 * 10**4 + 2  # total possible values in nums plus one dummy 
    tree = [0] * size 
    result = [] 
    for num in reversed(nums): 
        smaller_count = query(num + offset, tree) 
        result.append(smaller_count) 
        update(num + offset, 1, tree, size) 
    return reversed(result)