# python segment tree - count of smaller numbers after self

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)

• Could you please add some test cases? Commented Jan 6, 2020 at 12:36
• Note: you can probably just build the tree down-up and not up-down. This will remove the recursion and will be alot faster. Commented Jan 6, 2020 at 12:37

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