Count of Smaller Numbers After Self

Tested Python solution for LeetCode 315 with 22 pytest cases. Generate a practice environment with lcpy.

LeetCode 315, Hard. Topics: Array, Binary Search, Divide and Conquer, Binary Indexed Tree, Segment Tree, Merge Sort, Ordered Set, Treap. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 22 parametrized pytest cases, and a playground notebook:

lcpy gen -n 315   # by problem number
lcpy gen -s count_smaller_numbers_after_self   # by problem name

​

Problem

Given an integer array nums, return an integer array counts where counts[i] is the number of smaller elements to the right of nums[i].

​

Examples

Input: nums = [5,2,6,1]
Output: [2,1,1,0]

Explanation: To the right of 5 there are 2 smaller elements (2 and 1). To the right of 2 there is only 1 smaller element (1). To the right of 6 there is 1 smaller element (1). To the right of 1 there is 0 smaller element.

Input: nums = [-1]
Output: [0]

Input: nums = [-1,-1]
Output: [0,0]

​

Constraints

• 1 <= nums.length <= 10^5
• -10^4 <= nums[i] <= 10^4

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

class Solution:
    # Time: O(n log n)
    # Space: O(n)
    def count_smaller(self, nums: list[int]) -> list[int]:
        counts = [0] * len(nums)
        indices = list(range(len(nums)))

        def merge_sort(lo: int, hi: int) -> None:
            if hi - lo <= 1:
                return
            mid = (lo + hi) // 2
            merge_sort(lo, mid)
            merge_sort(mid, hi)

            merged: list[int] = []
            i, j = lo, mid
            while i < mid and j < hi:
                if nums[indices[j]] < nums[indices[i]]:
                    merged.append(indices[j])
                    j += 1
                else:
                    counts[indices[i]] += j - mid
                    merged.append(indices[i])
                    i += 1
            while i < mid:
                counts[indices[i]] += j - mid
                merged.append(indices[i])
                i += 1
            while j < hi:
                merged.append(indices[j])
                j += 1
            indices[lo:hi] = merged

        merge_sort(0, len(nums))
        return counts

​

Complexity

Time	Space
O(n log n)	O(n)

​

Tags