Skip to main content

Two Pointer

Partition Array According to Given Pivot Solution Code in C++/Python/Java/JS

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. If i < j and both elements are smaller (or larger) than pivot, then pi < pj.

Return nums after the rearrangement.

Example

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.

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.
💡
Try Solving "Partition Array According to Given Pivot" Yourself First !!

Figuring out the "Partition Array According to Given Pivot" requires pen and paper work. We need to code the Brute Force Solution and then observe the portions of the logic to be optimized. Let's figure it out together !!

Brute Force Approach

Intuition

As the question description says that we need to figure to arrange the elements based on few conditions. Let's first note them down.

1. We will be given an array of elements along with a pivot element.
2. We need to arrange the elements in such a way that, the numbers lesser than the pivot appears before it and the numbers greater than it appears after.
3. We need to make these changes in an array and return it.

The first intuitive thought that came to our mind is, let's first traverse the element from left to right.
We can maintain three lists where we could store the numbers lesser than the pivot, greater than the pivot or equal to the pivot element. Once, we are done with the traversal, we can simply merge them all into a single array and then return it.

Let's see how we can code it up !!

Approach

Initialize three lists:

  • less to store elements smaller than the pivot.
  • equal to store elements equal to the pivot.
  • greater to store elements greater than the pivot.

Iterate through the input array (nums):

  • If an element is smaller than pivot, add it to less.
  • If an element is equal to pivot, add it to equal.
  • If an element is greater than pivot, add it to greater.

Concatenate the lists in order:

  • Append equal to less.
  • Append greater to the updated less list.

Convert the list to an array:

  • Create a new array ans of the same size as nums.
  • Copy all elements from the merged list into ans.

Return the final array (ans).

Partition Array According to Given Pivot
Partition Array According to Given Pivot
Partition Array According to Given Pivot
partition-array-according-to-given-pivot-BruteForce

Dry-Run

Let's do a dry run of the algorithm using the input array:

Input:
nums = [9, 5, 8, 2, 1, 6, 5, 1, 10]
pivot = 5
Output:- [2,1,1,5,5,9,8,6,10]

Step 1: Initialize three lists

  • less = [] (for elements < pivot)
  • equal = [] (for elements == pivot)
  • greater = [] (for elements > pivot)

Step 2: Iterate through nums and categorize elements

  • 9 → greater = [9]
  • 5 → equal = [5]
  • 8 → greater = [9, 8]
  • 2 → less = [2]
  • 1 → less = [2, 1]
  • 6 → greater = [9, 8, 6]
  • 5 → equal = [5, 5]
  • 1 → less = [2, 1, 1]
  • 10 → greater = [9, 8, 6, 10]

Step 3: Merge the lists

less + equal + greater → [2, 1, 1] + [5, 5] + [9, 8, 6, 10]

Resulting array:
[2, 1, 1, 5, 5, 9, 8, 6, 10]

Final Output: [2, 1, 1, 5, 5, 9, 8, 6, 10]

Partition Array According to Given Pivot Solution

Now, let's checkout the Partition Array According to Given Pivot solution in c++, Java, Python and JavaScript

"Partition Array According to Given Pivot" Code in all Languages.
1. Partition Array According to Given Pivot solution in C++ Try on Compiler
#include <iostream>
#include <vector>
using namespace std;

class Solution {
public:
    // Function to partition the array based on pivot value
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        
        // Vector to store numbers less than pivot
        vector<int> less;

        // Vector to store numbers equal to pivot
        vector<int> equal;

        // Vector to store numbers greater than pivot
        vector<int> greater;

        // Iterate over each number in nums
        for (int num : nums) {
            
            // If the number is less than pivot, add to less vector
            if (num < pivot) {
                less.push_back(num);
            } 
            // If the number is greater than pivot, add to greater vector
            else if (num > pivot) {
                greater.push_back(num);
            } 
            // If the number is equal to pivot, add to equal vector
            else {
                equal.push_back(num);
            }
        }

        // Append equal elements to less vector
        less.insert(less.end(), equal.begin(), equal.end());

        // Append greater elements to less vector
        less.insert(less.end(), greater.begin(), greater.end());

        // Return the rearranged vector
        return less;
    }
};

// Main function to take input and test the function
int main() {
    int n;

    // Read the size of the array
    cin >> n;
    vector<int> nums(n);

    // Read array elements
    for (int i = 0; i < n; i++) {
        cin >> nums[i];
    }

    // Read pivot value
    int pivot;
    cin >> pivot;

    // Create an instance of Solution and get the result
    Solution solution;
    vector<int> result = solution.pivotArray(nums, pivot);

    // Print the output array
    for (int num : result) {
        cout << num << " ";
    }
    cout << endl;

    return 0;
}

2. Partition Array According to Given Pivot Solution in Java Try on Compiler
import java.util.*;

class Solution {

    // Method to partition the array based on pivot value
    public int[] pivotArray(int[] nums, int pivot) {
        
        // List to store numbers less than pivot
        List<Integer> less = new ArrayList<>();

        // List to store numbers equal to pivot
        List<Integer> equal = new ArrayList<>();

        // List to store numbers greater than pivot
        List<Integer> greater = new ArrayList<>();

        // Iterate over each number in nums
        for (int num : nums) {
            
            // If the number is less than pivot, add to less list
            if (num < pivot) {
                less.add(num);
            } 
            // If the number is greater than pivot, add to greater list
            else if (num > pivot) {
                greater.add(num);
            } 
            // If the number is equal to pivot, add to equal list
            else {
                equal.add(num);
            }
        }

        // Merge all three lists together
        less.addAll(equal);
        less.addAll(greater);

        // Initialize index for final answer array
        int i = 0;

        // Create result array with the same size as input
        int[] ans = new int[nums.length];

        // Copy values from merged list to array
        for (int num : less) {
            ans[i++] = num;
        }

        // Return the final rearranged array
        return ans;
    }

    // Main method to take input and test the function
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);

        // Read the size of the array
        int n = scanner.nextInt();
        int[] nums = new int[n];

        // Read array elements
        for (int i = 0; i < n; i++) {
            nums[i] = scanner.nextInt();
        }

        // Read pivot value
        int pivot = scanner.nextInt();
        scanner.close();

        // Create an instance of Solution and get the result
        Solution solution = new Solution();
        int[] result = solution.pivotArray(nums, pivot);

        // Print the output array
        System.out.println(Arrays.toString(result));
    }
}

3. Partition Array According to Given Pivot Solution in Python Try on Compiler
class Solution:
    # Function to partition the array based on pivot value
    def pivotArray(self, nums, pivot):
        
        # List to store numbers less than pivot
        less = []

        # List to store numbers equal to pivot
        equal = []

        # List to store numbers greater than pivot
        greater = []

        # Iterate over each number in nums
        for num in nums:
            
            # If the number is less than pivot, add to less list
            if num < pivot:
                less.append(num)

            # If the number is greater than pivot, add to greater list
            elif num > pivot:
                greater.append(num)

            # If the number is equal to pivot, add to equal list
            else:
                equal.append(num)

        # Combine all three lists and return the result
        return less + equal + greater


# Main function to take input and test the function
if __name__ == "__main__":
    
    # Read the size of the array
    n = int(input())
    
    # Read array elements
    nums = list(map(int, input().split()))
    
    # Read pivot value
    pivot = int(input())

    # Create an instance of Solution and get the result
    solution = Solution()
    result = solution.pivotArray(nums, pivot)

    # Print the output array
    print(result)

4. Partition Array According to Given Pivot solution in JavaScript Try on Compiler
const fs = require('fs');

// Function to partition the array based on pivot value
function pivotArray(nums, pivot) {
    
    // Array to store numbers less than pivot
    let less = [];

    // Array to store numbers equal to pivot
    let equal = [];

    // Array to store numbers greater than pivot
    let greater = [];

    // Iterate over each number in nums
    for (let num of nums) {
        
        // If the number is less than pivot, add to less array
        if (num < pivot) {
            less.push(num);
        } 
        // If the number is greater than pivot, add to greater array
        else if (num > pivot) {
            greater.push(num);
        } 
        // If the number is equal to pivot, add to equal array
        else {
            equal.push(num);
        }
    }

    // Combine all three arrays and return the result
    return [...less, ...equal, ...greater];
}

// Read input from file (sync)
let input = fs.readFileSync(0, 'utf-8').trim().split("\n");

// Read the size of the array
let n = parseInt(input[0]);

// Read array elements
let nums = input[1].split(" ").map(Number);

// Read pivot value
let pivot = parseInt(input[2]);

// Get the result
let result = pivotArray(nums, pivot);

// Print the output array
console.log(result);

Partition Array According to Given Pivot Complexity Analysis

Time Complexity: O(n)

Iterating through the array (nums):

    • Each element is compared to pivot and placed into one of the three lists (less, equal, or greater).
    • This takes O(N) time, where N is the length of nums.

Merging the lists: Appending equal and greater to less takes O(N) time.

Copying elements to the output array (ans): Copying from the merged list to ans also takes O(N) time.

Total time complexity:
O(N)+O(N)+O(N)=O(N)
Thus, the overall time complexity is O(N).

Space Complexity: O(n)

Auxiliary Space Complexity
Auxiliary space refers to the extra space used by the algorithm excluding the input and output space.

  • We create three additional lists: less, equal, and greater, which together hold N elements.
  • These lists take O(N) extra space.
  • The integer variables used (i, loop variables) take O(1) space.

Total auxiliary space complexity: O(N).

Total Space Complexity

  • Auxiliary space (extra space used apart from input).
  • Output space (the array returned).
  • The algorithm creates a new output array (ans) of size N.
  • Since the input array is read-only and not modified, we count the extra lists and the output array.

Total space used:

  • less, equal, greater: O(N)
  • Output array ans: O(N)
  • Other variables: O(1)

Thus, the total space complexity is:
O(N)+O(N)+O(1)=O(N).


The Brute Force approach for Partition Array According to Given Pivot invloved too much of extra space. In order to ace the interview process we need to figure out the space optimal solution as well. Let's gather few observations and finalise the optimal algorithm!!

Optimal Approach

Intuition

What do we need to do? We have an array and a special number called pivot. How should we arrange the numbers?
Smaller numbers go first.
Pivot numbers stay in the middle.
Bigger numbers go last.

But wait, does the order of elements matter? Yes! The relative order of numbers should not change.
Now, how do we think about solving this?

  • If we go from left to right, can we just put numbers smaller than pivot at the front?
  • If we always insert them at the next available slot in the front, won’t they naturally be in order?

Great! That means we can use a pointer (left) to track where the next smaller number should go.

How do we place greater numbers correctly?
Can we do something similar but in reverse? If we go from right to left and place numbers greater than pivot at the back, would they also be in order?

Instead of knowing where to start, can we just begin placing from the end of the array? So, we use another pointer (right) to track where the next greater number should go, filling the array backwards.

Wait, what about numbers equal to pivot? If we’re placing smaller numbers from the front and greater numbers from the back, what happens to the middle? Wouldn’t all the empty slots automatically be for pivot numbers?

So instead of tracking them separately, can we just fill the remaining spots with pivot in the end? This way, we don’t need to calculate boundaries—we just let them form naturally as we process the array!

Let's now check the algorithm

Partition Array According to Given Pivot Algorithm

  1. We implemented an approach to partition the array around a given pivot while maintaining relative order.
  2. First, we created an output array of the same length as the input. We used two pointers: lessI to track the position for elements smaller than the pivot and greaterI to track the position for elements greater than the pivot.
  3. By iterating from both ends of the array simultaneously, we placed smaller elements at the beginning and larger elements at the end. Finally, we filled the remaining positions with the pivot value. This approach ensures an efficient in-place partitioning of the array.

Let's see how we can code the Partition Array According to Given Pivot Solution.

Approach

Initialize an output array (ans) of the same length as nums.

Set two pointers:

  • lessI = 0 (tracks the position for elements smaller than the pivot).
  • greaterI = nums.length - 1 (tracks the position for elements greater than the pivot).

Iterate through nums using two indices (i from the start, j from the end):

  • If nums[i] < pivot, place it at ans[lessI] and increment lessI.
  • If nums[j] > pivot, place it at ans[greaterI] and decrement greaterI.

Fill the remaining positions between lessI and greaterI with pivot.

Return the output array (ans).

Let us understand this with a video ,

0:00
/2:26

partition-array-according-to-given-pivot-optimal animation

Dry-Run

Input:
nums = [9, 5, 8, 2, 1, 6, 5, 1, 10]
pivot = 5
Output:- [2,1,1,5,5,9,8,6,10]

Step 1: Initialize Variables

  • Create an array ans of size nums.length.
  • Set lessI = 0 to track positions for elements < pivot.
  • Set greaterI = nums.length - 1 to track positions for elements > pivot.

Step 2: First Loop - Processing nums[i] and nums[j]

  • i = 0, j = 8
    • nums[0] = 9 (not less than pivot, no action).
    • nums[8] = 10 (greater than pivot), so place 10 at ans[8].
    • Decrement greaterI to 7.
  • i = 1, j = 7
    • nums[1] = 5 (equal to pivot, no action).
    • nums[7] = 1 (less than pivot), so place 1 at ans[0].
    • Increment lessI to 1.
  • i = 2, j = 6
    • nums[2] = 8 (not less than pivot, no action).
    • nums[6] = 5 (equal to pivot, no action).
  • i = 3, j = 5
    • nums[3] = 2 (less than pivot), so place 2 at ans[1].
    • Increment lessI to 2.
    • nums[5] = 6 (greater than pivot), so place 6 at ans[7].
    • Decrement greaterI to 6.
  • i = 4, j = 4
    • nums[4] = 1 (less than pivot), so place 1 at ans[2].
    • Increment lessI to 3.
  • i = 5, j = 3
    • nums[5] = 6 (already placed in ans, no action).
    • nums[3] = 2 (already placed in ans, no action).
  • i = 6, j = 2
    • nums[6] = 5 (equal to pivot, no action).
    • nums[2] = 8 (greater than pivot), so place 8 at ans[6].
    • Decrement greaterI to 5.
  • i = 7, j = 1
    • nums[7] = 1 (already placed in ans, no action).
    • nums[1] = 5 (equal to pivot, no action).
  • i = 8, j = 0
    • nums[8] = 10 (already placed in ans, no action).
    • nums[0] = 9 (greater than pivot), so place 9 at ans[5].
    • Decrement greaterI to 4.

Step 3: Fill Remaining Positions with Pivot

  • Fill ans[3] to ans[4] with pivot = 5.

Final Output

  • ans = [2, 1, 1, 5, 5, 9, 8, 6, 10]

Partition Array According to Given Pivot Solution

Now, let's checkout the Partition Array According to Given Pivot solution in c++, Java, Python and JavaScript

"Partition Array According to Given Pivot" Code in all Languages.
1. Partition Array According to Given Pivot solution in C++ Try on Compiler
#include <iostream>
#include <vector>

using namespace std;

vector<int> pivotArray(vector<int>& nums, int pivot) {
    // Create a result vector with the same size as input
    vector<int> ans(nums.size());

    // Initialize pointers for elements less than and greater than pivot
    int lessI = 0;
    int greaterI = nums.size() - 1;

    // Traverse the array from both ends
    for (int i = 0, j = nums.size() - 1; i < nums.size(); i++, j--) {
        // If current element is less than pivot, place it at lessI index
        if (nums[i] < pivot) {
            ans[lessI] = nums[i];
            lessI++;
        }

        // If current element is greater than pivot, place it at greaterI index
        if (nums[j] > pivot) {
            ans[greaterI] = nums[j];
            greaterI--;
        }
    }

    // Fill the remaining positions with pivot
    while (lessI <= greaterI) {
        ans[lessI] = pivot;
        lessI++;
    }

    return ans;
}

int main() {
    // Read input from user
    int n;
    cout << "Enter number of elements: ";
    cin >> n;

    vector<int> nums(n);

    cout << "Enter array elements: ";
    for (int i = 0; i < n; i++) {
        cin >> nums[i];
    }

    int pivot;
    cout << "Enter pivot: ";
    cin >> pivot;

    // Call the function
    vector<int> result = pivotArray(nums, pivot);

    // Print the output array
    cout << "Rearranged Array: ";
    for (int num : result) {
        cout << num << " ";
    }
    cout << endl;

    return 0;
}

2. Partition Array According to Given Pivot Solution in Java Try on Compiler
import java.util.*;

class Solution {

    public int[] pivotArray(int[] nums, int pivot) {
        // Create an array to store the result
        int[] ans = new int[nums.length];

        // Initialize pointers for elements less than and greater than pivot
        int lessI = 0;
        int greaterI = nums.length - 1;

        // Traverse the array from both ends
        for (int i = 0, j = nums.length - 1; i < nums.length; i++, j--) {
            // If current element is less than pivot, place it at lessI index
            if (nums[i] < pivot) {
                ans[lessI] = nums[i];
                lessI++;
            }

            // If current element is greater than pivot, place it at greaterI index
            if (nums[j] > pivot) {
                ans[greaterI] = nums[j];
                greaterI--;
            }
        }

        // Fill the remaining positions with pivot
        while (lessI <= greaterI) {
            ans[lessI] = pivot;
            lessI++;
        }

        return ans;
    }

    public static void main(String[] args) {
        // Scanner to take input from the user
        Scanner sc = new Scanner(System.in);
        
        // Read the size of the array
        System.out.print("Enter number of elements: ");
        int n = sc.nextInt();

        // Create an array to store user input
        int[] nums = new int[n];
        
        // Read array elements
        System.out.println("Enter array elements: ");
        for (int i = 0; i < n; i++) {
            nums[i] = sc.nextInt();
        }
        
        // Read pivot value
        System.out.print("Enter pivot: ");
        int pivot = sc.nextInt();
        
        // Call pivotArray function
        Solution sol = new Solution();
        int[] result = sol.pivotArray(nums, pivot);

        // Print the output array
        System.out.println("Rearranged Array: " + Arrays.toString(result));
    }
}

3. Partition Array According to Given Pivot Solution in Python Try on Compiler
def pivot_array(nums, pivot):
    # Create a result list with the same size as input
    ans = [0] * len(nums)

    # Initialize pointers for elements less than and greater than pivot
    less_i = 0
    greater_i = len(nums) - 1

    # Traverse the array from both ends
    for i, j in zip(range(len(nums)), reversed(range(len(nums)))):
        # If current element is less than pivot, place it at less_i index
        if nums[i] < pivot:
            ans[less_i] = nums[i]
            less_i += 1

        # If current element is greater than pivot, place it at greater_i index
        if nums[j] > pivot:
            ans[greater_i] = nums[j]
            greater_i -= 1

    # Fill the remaining positions with pivot
    while less_i <= greater_i:
        ans[less_i] = pivot
        less_i += 1

    return ans


if __name__ == "__main__":
    # Read input from user
    n = int(input("Enter number of elements: "))
    nums = list(map(int, input("Enter array elements: ").split()))
    pivot = int(input("Enter pivot: "))

    # Call the function and print result
    result = pivot_array(nums, pivot)
    print("Rearranged Array:", result)

4. Partition Array According to Given Pivot solution in JavaScript Try on Compiler
const fs = require("fs");

// Function to rearrange elements based on pivot
function pivotArray(nums, pivot) {
    // Create a result array with the same size as input
    let ans = new Array(nums.length);

    // Initialize pointers for elements less than and greater than pivot
    let lessI = 0;
    let greaterI = nums.length - 1;

    // Traverse the array from both ends
    for (let i = 0, j = nums.length - 1; i < nums.length; i++, j--) {
        // If current element is less than pivot, place it at lessI index
        if (nums[i] < pivot) {
            ans[lessI] = nums[i];
            lessI++;
        }

        // If current element is greater than pivot, place it at greaterI index
        if (nums[j] > pivot) {
            ans[greaterI] = nums[j];
            greaterI--;
        }
    }

    // Fill the remaining positions with pivot
    while (lessI <= greaterI) {
        ans[lessI] = pivot;
        lessI++;
    }

    return ans;
}

// Read input from a file
const input = fs.readFileSync("input.txt", "utf8").trim().split("\n");
const nums = input[1].split(" ").map(Number);
const pivot = Number(input[2]);

// Call the function and print result
const result = pivotArray(nums, pivot);
console.log("Rearranged Array:", result.join(" "));

Partition Array According to Given Pivot Complexity Analysis

Time Complexity: O(n)

The algorithm consists of the following steps:

Initialize the output array (ans)O(N)

Iterate through the array (nums) once with two indices (i and j):
Each element is checked and placed in ans if it's < pivot or > pivot. This takes O(N) time.

Fill the remaining positions with pivotO(N) in the worst case (if most elements are equal to pivot).

Total time complexity: O(N)+O(N)+O(N)=O(N)
Thus, the overall time complexity is O(N).

Space Complexity: O(n)

Auxiliary Space Complexity
Auxiliary space refers to the extra space used by the algorithm excluding the input and output space.

  • The algorithm only uses a few integer variables (lessI, greaterI, i, j), which take O(1) space.
  • No additional data structures (like lists or hash maps) are used.

Total auxiliary space complexity: O(1).

Total Space Complexity
Total space includes both:

  • Auxiliary space (extra space used apart from input).
  • Output space (the array returned).
  • The algorithm creates a new output array (ans) of size N, which takes O(N) space.
  • Since the input array is read-only and not modified, we only count ans.

Total space complexity: O(N)


Similar Problems

Partition List

Rearrange Array Elements by Sign