Partition Array According to Given Pivot

Med

#1968Time: O(N), where N is the number of elements in `nums`. We iterate through the array three times, so the total time is O(N) + O(N) + O(N) = O(N).Space: O(N), where N is the number of elements in `nums`. We use an auxiliary array of size N to store the result. If the output array is not considered extra space, the complexity would be O(1), but typically it is.

Patterns

Two Pointers

Data structures

Array

Prompt

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 p_i, p_j where p_i is the new position of the i^th element and p_j is the new position of the j^th element. If i < j and both elements are smaller (or larger) than pivot, then p_i < p_j.

Return nums after the rearrangement.

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 <= 10⁵
-10⁶ <= nums[i] <= 10⁶
pivot equals to an element of nums.

Approaches

2 approaches with complexity analysis and trade-offs.

This is a straightforward approach where we iterate through the input array three separate times. Each pass is responsible for collecting elements belonging to one category (less than, equal to, or greater than the pivot). We use an auxiliary array to build the result. By iterating through the input array in its original order for each category, we naturally preserve the relative ordering of the elements.

Algorithm

Create a new result array ans of the same size as nums.
Initialize an index i = 0 for the ans array.
First pass: Iterate through nums. For each num < pivot, add it to ans[i] and increment i.
Second pass: Iterate through nums. For each num == pivot, add it to ans[i] and increment i.
Third pass: Iterate through nums. For each num > pivot, add it to ans[i] and increment i.
Return ans.

Walkthrough

The core idea is to create a new array and fill it sequentially. We first fill it with all elements smaller than the pivot, then with all elements equal to the pivot, and finally with all elements greater than the pivot.

The algorithm proceeds as follows:

Create a new integer array result with the same length as nums.
Initialize a pointer, index, to 0. This pointer will keep track of the next available position in the result array.
First Pass (Less than pivot): Iterate through nums from left to right. If an element num is less than pivot, place it at result[index] and increment index.
Second Pass (Equal to pivot): Iterate through nums again. If an element num is equal to pivot, place it at result[index] and increment index.
Third Pass (Greater than pivot): Iterate through nums a final time. If an element num is greater than pivot, place it at result[index] and increment index.
After these three passes, the result array will contain the partitioned elements in the correct, stable order. Return result.

Here is a code snippet demonstrating this approach:

1class Solution {2    public int[] partitionArray(int[] nums, int pivot) {3        int n = nums.length;4        int[] result = new int[n];5        int index = 0;6 7        // 1. Pass for elements less than pivot8        for (int num : nums) {9            if (num < pivot) {10                result[index++] = num;11            }12        }13 14        // 2. Pass for elements equal to pivot15        for (int num : nums) {16            if (num == pivot) {17                result[index++] = num;18            }19        }20 21        // 3. Pass for elements greater than pivot22        for (int num : nums) {23            if (num > pivot) {24                result[index++] = num;25            }26        }27 28        return result;29    }30}

Complexity

Time

O(N), where N is the number of elements in `nums`. We iterate through the array three times, so the total time is O(N) + O(N) + O(N) = O(N).

Space

O(N), where N is the number of elements in `nums`. We use an auxiliary array of size N to store the result. If the output array is not considered extra space, the complexity would be O(1), but typically it is.

Trade-offs

Pros

Simple and intuitive to understand and implement.
Correctly preserves the relative order of elements within the 'less than' and 'greater than' partitions.

Cons

Requires three full traversals of the input array, which is less efficient than a single or two-pass solution in terms of constant factors.
Requires extra space proportional to the input size.

This approach improves upon the three-pass method by reducing the number of traversals to two. In the first pass, we count the number of elements that are less than the pivot and equal to the pivot. These counts allow us to calculate the exact boundaries for the three partitions in the final array. The second pass then iterates through the input array once more, placing each element directly into its correct spot in the result array.

Algorithm

Initialize lessCount = 0 and equalCount = 0.
First pass: Iterate through nums to count lessCount (elements < pivot) and equalCount (elements == pivot).
Define starting indices: lessIndex = 0, equalIndex = lessCount, greaterIndex = lessCount + equalCount.
Create a new result array ans.
Second pass: Iterate through nums and place each element num into ans at the correct partition index (lessIndex, equalIndex, or greaterIndex) and increment the corresponding index.

Walkthrough

This method is more optimized as it avoids redundant iterations.

The algorithm is as follows:

First Pass (Counting): Traverse the nums array once to count the number of elements smaller than the pivot (lessCount) and the number of elements equal to the pivot (equalCount).
Determine Partition Boundaries: Use the counts to determine the starting indices for each partition in the result array.
- The 'less than' partition starts at index 0.
- The 'equal to' partition starts at index lessCount.
- The 'greater than' partition starts at index lessCount + equalCount.
Second Pass (Placing): Create a new result array. Initialize three pointers: lessIndex = 0, equalIndex = lessCount, and greaterIndex = lessCount + equalCount. Iterate through the original nums array. For each element:
- If it's less than pivot, place it at result[lessIndex] and increment lessIndex.
- If it's equal to pivot, place it at result[equalIndex] and increment equalIndex.
- If it's greater than pivot, place it at result[greaterIndex] and increment greaterIndex.
Return the result array. This single placement pass ensures relative order is maintained.

Here is a code snippet for this approach:

1class Solution {2    public int[] partitionArray(int[] nums, int pivot) {3        int n = nums.length;4        int[] result = new int[n];5        6        int lessCount = 0;7        int equalCount = 0;8        9        // 1. First pass: Count elements10        for (int num : nums) {11            if (num < pivot) {12                lessCount++;13            } else if (num == pivot) {14                equalCount++;15            }16        }17        18        int lessIndex = 0;19        int equalIndex = lessCount;20        int greaterIndex = lessCount + equalCount;21        22        // 2. Second pass: Place elements23        for (int num : nums) {24            if (num < pivot) {25                result[lessIndex++] = num;26            } else if (num == pivot) {27                result[equalIndex++] = num;28            } else {29                result[greaterIndex++] = num;30            }31        }32        33        return result;34    }35}

Complexity

Time

O(N), where N is the number of elements in `nums`. We iterate through the array twice, once for counting and once for placing. The total time is O(N) + O(N) = O(N).

Space

O(N), where N is the number of elements in `nums`. We use an auxiliary array of size N to store the result.

Trade-offs

Pros

More efficient than the three-pass approach in terms of constant factors (2N vs 3N operations).
The logic is clean and directly places elements, preserving relative order.
Still relatively simple to implement.

Cons

Still requires O(N) extra space, which might be a concern for very large inputs if memory is constrained.

Solutions

Solution

1class Solution { public int [] pivotArray ( int [] nums , int pivot ) { int n = nums . length ; int [] ans = new int [ n ]; int k = 0 ; for ( int x : nums ) { if ( x < pivot ) { ans [ k ++] = x ; } } for ( int x : nums ) { if ( x == pivot ) { ans [ k ++] = x ; } } for ( int x : nums ) { if ( x > pivot ) { ans [ k ++] = x ; } } return ans ; } }

1class Solution { public: vector < int > pivotArray ( vector < int >& nums , int pivot ) { vector < int > ans ; for ( int & x : nums ) if ( x < pivot ) ans . push_back ( x ); for ( int & x : nums ) if ( x == pivot ) ans . push_back ( x ); for ( int & x : nums ) if ( x > pivot ) ans . push_back ( x ); return ans ; } };

1class Solution : def pivotArray ( self , nums : List [ int ], pivot : int ) -> List [ int ]: a , b , c = [], [], [] for x in nums : if x < pivot : a . append ( x ) elif x == pivot : b . append ( x ) else : c . append ( x ) return a + b + c

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem