Problem Statement in English

You’re given a 0-indexed integer array nums and an integer pivot. Rearrange nums such that the following conditions are satisfied:

  • Every element less than pivot appears before every element greater than pivot.
  • Every element equal to pivot appears in between the elements less than and greater than pivot.
  • The relative order of the elements less than pivot and the elements greater than pivot is maintained.

Return nums after the rearrangement.

Approach

Since we need to maintain the relative order of the elements less than and greater than pivot, we can use three separate lists to store the elements less than, equal to, and greater than pivot. Finally, we can concatenate these lists to get the desired result.

An optimisation would be to simply count the number of occurrences of pivot and then construct the final list using list concatenation - so that’s just 2 lists.

And we’re done!

Solution in Python


class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        lt = []
        gt = []
        c = 0

        for n in nums:
            if n < pivot:
                lt.append(n)
            elif n == pivot:
                c += 1
            else:
                gt.append(n)

        return lt + [pivot] * c + gt

Complexity

  • Time: $O(n)$
    Since we need to iterate through the entire list once to separate the elements, the time complexity is linear with respect to the size of the input list nums.

  • Space: $O(n)$
    Since we are using additional lists to store the elements less than, equal to, and greater than pivot, the space complexity is also linear with respect to the size of the input list nums.

And we are done.