Partition Array According to Given Pivot

MEDIUM

Description

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

Checkout 2 different approaches to solve Partition Array According to Given Pivot. Click on different approaches to view the approach and algorithm in detail.

Three Passes with Extra Space

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.

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:

class Solution {
    public int[] partitionArray(int[] nums, int pivot) {
        int n = nums.length;
        int[] result = new int[n];
        int index = 0;

        // 1. Pass for elements less than pivot
        for (int num : nums) {
            if (num < pivot) {
                result[index++] = num;
            }
        }

        // 2. Pass for elements equal to pivot
        for (int num : nums) {
            if (num == pivot) {
                result[index++] = num;
            }
        }

        // 3. Pass for elements greater than pivot
        for (int num : nums) {
            if (num > pivot) {
                result[index++] = num;
            }
        }

        return result;
    }
}

Complexity Analysis

Time Complexity: 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 Complexity: 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.

Pros and Cons

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.

Code Solutions

Checking out 3 solutions in different languages for Partition Array According to Given Pivot. Click on different languages to view the code.

class 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 ; } }

Video Solution

Watch the video walkthrough for Partition Array According to Given Pivot

Patterns:

Two Pointers

Data Structures:

Array