2
\$\begingroup\$

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)

\$\endgroup\$
  • \$\begingroup\$ Could you please add some test cases? \$\endgroup\$ – Yonlif Jan 6 at 12:36
  • \$\begingroup\$ Note: you can probably just build the tree down-up and not up-down. This will remove the recursion and will be alot faster. \$\endgroup\$ – Yonlif Jan 6 at 12:37
1
\$\begingroup\$

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

| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.