LeetCode 315. Count of Smaller Numbers After Self

MergeSort Python

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        def sort(l, r):
            if l >= r: return
            mid = (l+r) // 2
            sort(l, mid)
            sort(mid+1, r)
            i = l
            j = mid+1
            tmp = []
            add = 0
            while i <= mid and j <= r:
                if store[j][0] < store[i][0]:
                    tmp.append(store[j])
                    add += 1
                    j += 1
                else:
                    tmp.append(store[i])
                    res[store[i][1]] += add
                    i += 1
            while i <= mid:
                tmp.append(store[i])
                res[store[i][1]] += add
                i += 1
            while j <= r:
                tmp.append(store[j])
                j += 1
            i = 0
            for k in range(l, r+1):
                store[k] = tmp[i]
                i += 1
        
        n = len(nums)
        res = [0]*n
        
        store = []
        for i in range(n):
            store.append((nums[i], i))
        
        sort(0, n-1)
        return res
        
        

Binary Indexed Tree

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        N = 20001
        tr = [0]*N
        
        def lowbit(x):
            return x & -x
        def add(x):
            nonlocal N
            i = x
            while i < N:
                tr[i] += 1
                i += lowbit(i)
        def query(x):
            i = x
            count = 0
            while i > 0:
                count += tr[i]
                i -= lowbit(i)
            return count
        
        n = len(nums)
        res = [0]*n
        for i in range(n-1, -1, -1):
            x = nums[i] + 10001
            add(x)
            res[i] = query(x-1)
        return res