Count of Smaller Numbers After Self LeetCode solution in Python

1
# Runtime: 6019 ms (Top 15.52%) | Memory: 35.5 MB (Top 33.31%)
2
class Solution:
3
    def countSmaller(self, nums: List[int]) -> List[int]:
4
        # build the binary indexed tree
5
        num_buckets = (
6
            10**4 + 10**4 + 1
7
        )  # 10**4 negative + 10**4 positive numbers + bucket at 0
8
        tree = [0] * (
9
            num_buckets + 1
10
        )  # add 1 because binary indexed tree data starts at index 1
11

12
        result = [0] * len(nums)
13

14
        # iterate from right to left
15
        for result_index in range(len(nums) - 1, -1, -1):
16
            n = nums[result_index]
17
            # add 10**4 to n to account for negative numbers
18
            i = n + 10**4
19

20
            # convert to 1-based index for the tree
21
            i += 1
22
            # perform range sum query of buckets [-inf, n-1], where n is current number
23
            # because we want n - 1 for range sum query of [-inf, n-1], not n, subtract 1 from i:
24
            i -= 1
25

26
            val = 0
27
            while i != 0:
28
                val += tree[i]
29
                # get parent node by subtracting least significant set bit
30
                i -= i & -i
31

32
            result[result_index] = val
33

34
            # update the binary indexed tree with new bucket
35
            i = n + 10**4
36
            i += 1
37
            while i < len(tree):
38
                tree[i] += 1
39
                # get next node to update by adding the least significant set bit
40
                i += i & -i
41

42
        return result