Partition Array According to Given Pivot

Problem

You are 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.
  • More formally, consider every pi, pj where pi is the new position of the ith element and pj is the new position of the jth element. For elements less than pivot, if i < j and nums[i] < pivot and nums[j] < pivot, then pi < pj. Similarly for elements greater than pivot, if i < j and nums[i] > pivot and nums[j] > pivot, then pi < pj.

Return nums after the rearrangement.

Examples

Example 1:

Input: nums = [9,12,5,10,14,3,10], pivot = 10
Output: [9,5,3,10,10,12,14]
Explanation: 
The elements 9, 5, and 3 are less than the pivot so they are on the left side of the array.
The elements 12 and 14 are greater than the pivot so they are on the right side of the array.
The relative ordering of the elements less than and greater than pivot is also maintained. [9, 5, 3] and [12, 14] are the respective orderings.

Example 2:

Input: nums = [-3,4,3,2], pivot = 2
Output: [-3,2,4,3]
Explanation: 
The element -3 is less than the pivot so it is on the left side of the array.
The elements 4 and 3 are greater than the pivot so they are on the right side of the array.
The relative ordering of the elements less than and greater than pivot is also maintained. [-3] and [4, 3] are the respective orderings.

Constraints:

• 1 <= nums.length <= 105
• -106 <= nums[i] <= 106
• pivot equals to an element of nums.

Solution

Method 1 - Creating 3 lists and update the arrays

To rearrange the elements of the array based on the given pivot, we need to maintain three separate lists:

• One for elements less than the pivot.
• One for elements equal to the pivot.
• One for elements greater than the pivot.

By iterating through the array and appending each element to the appropriate list, we effectively partition the array into the required format while maintaining relative order. Finally, we concatenate these three lists to obtain the rearranged array.

Approach

1. Initialize Three Lists:
   • One for elements less than pivot.
   • One for elements equal to pivot.
   • One for elements greater than pivot.
2. Iterate Through the Original List:
   • Append each element to the corresponding list based on its comparison with pivot.
3. Concatenate the Three Lists:
   • Combine all three lists to get the partitioned list.
4. Return the Partitioned List:
   • Return the concatenated list as the result.

Code

Java

public class Solution {
    public int[] pivotArray(int[] nums, int pivot) {
        List<Integer> less = new ArrayList<>();
        List<Integer> equal = new ArrayList<>();
        List<Integer> greater = new ArrayList<>();

        for (int num : nums) {
            if (num < pivot) {
                less.add(num);
            } else if (num == pivot) {
                equal.add(num);
            } else {
                greater.add(num);
            }
        }

        int[] ans = new int[nums.length];
        int index = 0;

        for (int num : less) {
            ans[index++] = num;
        }
        for (int num : equal) {
            ans[index++] = num;
        }
        for (int num : greater) {
            ans[index++] = num;
        }

        return ans;
    }
}

Python

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        less_than_pivot: List[int] = []
        equal_to_pivot: List[int] = []
        greater_than_pivot: List[int] = []
        
        for num in nums:
            if num < pivot:
                less_than_pivot.append(num)
            elif num == pivot:
                equal_to_pivot.append(num)
            else:
                greater_than_pivot.append(num)
        
        return less_than_pivot + equal_to_pivot + greater_than_pivot

Complexity

• ⏰ Time complexity: O(n), where n is the number of elements in nums, as we are iterating through the list once.
• 🧺 Space complexity: O(n), where n is the number of elements in nums, for storing elements in the three separate lists.

Method 2 - Using Two Pointer Approach

We can also use two pointer technique:

• Using the two-pointer approach, we process the nums array in one traversal while maintaining relative order.
  • lessIdx to place elements less than pivot at the beginning of the array.
  • greaterIdx to place elements greater than pivot from the end of the array.
• Create an auxiliary array ans of the same length as nums to store the rearranged numbers.

Here is the algorithm:

• Numbers < pivot get placed sequentially via lessIdx.
• Numbers > pivot get stored via greaterIdx
• All occurrences of pivot fill the middle (while loop) once the pointers converge.

Code

Java

class Solution {

    public int[] pivotArray(int[] nums, int pivot) {
        int[] ans = new int[nums.length];
        int lessIdx = 0; // Pointer for < pivot
        int greaterIdx = nums.length - 1; // Pointer for > pivot

        for (int i = 0, j = nums.length - 1; i < nums.length; i++, j--) {
            if (nums[i] < pivot) {
                ans[lessIdx] = nums[i];
                lessIdx++;
            }
            if (nums[j] > pivot) {
                ans[greaterIdx] = nums[j];
                greaterIdx--;
            }
        }

        while (lessIdx <= greaterIdx) {
            ans[lessIdx] = pivot;
            lessIdx++;
        }

        return ans;
    }
}

Python

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        ans = [0] * len(nums)  # Result array
        lessIdx = 0  # Pointer for < pivot
        greaterIdx = len(nums) - 1  # Pointer for > pivot

        i, j = 0, len(nums) - 1
        while i < len(nums):
            if nums[i] < pivot:
                ans[lessIdx] = nums[i]
                lessIdx += 1
            if nums[j] > pivot:
                ans[greaterIdx] = nums[j]
                greaterIdx -= 1
            i += 1
            j -= 1

        # Fill all the pivot values in the middle
        while lessIdx <= greaterIdx:
            ans[lessIdx] = pivot
            lessIdx += 1

        return ans

Complexity

• ⏰ Time complexity: O(n) as we traverse the array three times with linear time operations.
• 🧺 Space complexity: O(n) due to the use of the auxiliary array ans.