python segment tree - count of smaller numbers after self

Question

I am trying to solve this question https://leetcode.com/problems/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].

My code:

import bisect

from functools import lru_cache


def merge(a, b):
    i, j = 0, 0
    res = []
    while i < len(a) and j < len(b):
        if a[i] < b[j]:
            res.append(a[i])
            i += 1
        else:
            res.append(b[j])
            j += 1
    while i < len(a):
        res.append(a[i])
        i += 1
    while j < len(b):
        res.append(b[j])
        j += 1
    return res


class SegmentTree:
    def __init__(self, nums):
        self.nums = nums
        self.tree = [0] * (len(nums)*4)
        if len(nums)>0:
            self._build()

    def _build(self):
        def build(l, r, index):
            if l == r:
                self.tree[index] = [self.nums[l]]
                return self.tree[index]
            else:
                mid = (l+r)//2
                left_set = build(l, mid, index*2)
                right_set = build(mid+1, r, index*2+1)
                m = merge(left_set, right_set)
                self.tree[index] = m
                return m
        build(0, len(self.nums)-1, 1)

    def get_range(self, l, r):

        @lru_cache(maxsize=None)
        def get_range(left_boundary, right_boudndary, index, l, r):
            if l > r:
                return []
            if left_boundary == right_boudndary:
                return self.tree[index]
            if left_boundary == l and right_boudndary == r:
                return self.tree[index]
            else:
                mid = (left_boundary + right_boudndary)//2
                left = get_range(left_boundary, mid, index*2, l, min(r, mid))
                right = get_range(mid+1, right_boudndary,
                                index*2+1, max(l, mid+1), r)
                return merge(left, right)

        return get_range(0, len(self.nums)-1, 1, l, r)


class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        s = SegmentTree(nums)
        result = []
        for i in range(len(nums)):
            res  = s.get_range(i+1, len(nums)-1)
            ans = bisect.bisect(res, nums[i]-1)
            result.append(ans)
        return result

Sadly this times out on the last case :(

How do I speed this up?

(To be clear: I am looking for a faster solution with segment trees, not fenwick tree)

Answer 1

So, the time complexity of the solution posted in the question is O(n^2 log(n)).

Answering each query takes n log(n), and we have n queries in total.

We don't necessarily need to merge the left and right subtrees to find the inversion count; given that the sublists are sorted, we can exploit binary search.

import bisect

from functools import lru_cache

def merge(left, right):
    res = []
    i, j = 0, 0 
    while i<len(left) and j<len(right):
        if left[i]<right[j]:
            res.append(left[i])
            i+=1 
        else:
            res.append(right[j])
            j+=1 

    while i<len(left):
        res.append(left[i])
        i+=1 
    while j<len(right):
        res.append(right[j])
        j+=1 
    return res


class SegmentTree:
    def __init__(self, nums):
        self.nums = nums 
        self.tree = {}
        if len(nums)>0:
            self._build(0, len(nums)-1, 1)
    def _build(self, l, r, index):
        if l==r:
            self.tree[index] = [self.nums[r]]
            return self.tree[index]
        else:
            mid = (l+r)//2 
            left= self._build(l, mid, index*2)
            right = self._build(mid+1, r, index*2+1)
            self.tree[index] = merge(left, right)
            return self.tree[index]

    def get_range(self, l,r, target):
        def get_range(left_boundary, right_boundary, l, r, index):
            if l>r:
                return 0 
            if left_boundary == right_boundary or (left_boundary == l and right_boundary==r):
                return bisect.bisect(self.tree[index], target)
            else:
                mid = (left_boundary+ right_boundary)//2 
                left = get_range(left_boundary, mid, l, min(r, mid), index*2)
                right = get_range(mid+1, right_boundary, max(l, mid+1), r, index*2+1)
                return left+right
        return get_range(0, len(self.nums)-1, l, r, 1)

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        s = SegmentTree(nums)
        result = []
        for i in range(len(nums)):
            res = s.get_range(i+1, len(nums)-1, nums[i]-1)
            result.append(res)
        return result