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)

• Yonlif: Could you please add some test cases?
• Yonlif: Note: you can probably just build the tree down-up and not up-down. This will remove the recursion and will be alot faster.

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