Problem Informations

• Problem Index: 2161
• Problem Link: Partition Array According to Given Pivot
• Topics: Array, Two Pointers, Simulation

Problem 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 \leq \text{nums.length} \leq 10^5$
• $-10^6 \leq \text{nums}[i] \leq 10^6$
• pivot equals to an element of nums.

Intuition

The problem essentially revolves around rearranging an array based on its elements in relation to a given pivot. The core idea is to divide the array into three distinct sections: elements less than the pivot, elements equal to the pivot, and elements greater than the pivot. By maintaining these sections and not altering the relative order of elements within them, we ensure that the array adheres to the given requirements, while preserving stability for elements with respect to the pivot.

Approach

The approach involves utilizing auxiliary data structures to categorize and reorder the elements efficiently:

1. Initialization:

   • Utilize two queues: less for storing elements less than the pivot, and greater for elements greater than the pivot. Queues are ideal here as they preserve the order of elements, which is essential for maintaining the relative order.

2. Classification:

   • Iterate through the given nums array. For every element in nums, compare it to the pivot.
     • If the element is less than pivot, add it to the less queue.
     • If the element is greater than pivot, add it to the greater queue.
     • Elements equal to pivot are not added to any queue as they will be dealt with later in-place.

3. Reconstruction:

   • Initialize an index idx at 0, which will be used to overwrite the nums array.
   • Dequeue all elements from the less queue and overwrite nums from the start, moving idx forward.
   • Next, fill the array from the current idx position with the pivot value itself. The number of pivot elements to place is derived from the array’s total size minus the combined size of less and greater queues.
   • Finally, dequeue all elements from the greater queue and continue to overwrite nums starting from the updated idx.

This method ensures that the original order within elements less than or greater than the pivot is maintained, while effectively partitioning the elements around the pivot. The time complexity of this approach is $O(n)$, where $n$ is the number of elements in nums, as each element is processed once, and operations related to queues are of constant time complexity.

Code

C++

class Solution {
public:
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        // Queues to store elements less than and greater than pivot
        queue<int> less, greater;

        // Classify elements based on their relationship to pivot
        for (int num : nums) {
            if (num < pivot) {
                less.push(num);
            } else if (num > pivot) {
                greater.push(num);
            }
        }

        int idx = 0;

        // Place elements less than pivot at the start of nums
        while (!less.empty()) {
            nums[idx++] = less.front();
            less.pop();
        }

        // Place elements equal to pivot next in nums
        for (; idx < nums.size() - greater.size(); idx++) {
            nums[idx] = pivot;
        }

        // Place elements greater than pivot at the end of nums
        while (!greater.empty()) {
            nums[idx++] = greater.front();
            greater.pop();
        }

        return nums;
    }
};

Python

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        # Queues to store elements less than and greater than pivot
        less = deque()
        greater = deque()

        # Classify elements based on their relationship to pivot
        for num in nums:
            if num < pivot:
                less.append(num)
            elif num > pivot:
                greater.append(num)

        idx = 0

        # Place elements less than pivot at the start of nums
        while less:
            nums[idx] = less.popleft()
            idx += 1

        # Place elements equal to pivot next in nums
        for _ in range(idx, len(nums) - len(greater)):
            nums[idx] = pivot
            idx += 1

        # Place elements greater than pivot at the end of nums
        while greater:
            nums[idx] = greater.popleft()
            idx += 1

        return nums

Complexity

• Time complexity: The time complexity of the solution is $O(n)$, where $n$ is the length of the input array nums. This is because the algorithm iterates through the nums array once to classify the elements into “less” and “greater” queues, and then iterates through these queues to rearrange the elements in nums. Each operation (pushing to or popping from a queue, and writing back to nums) takes constant time, $O(1)$, resulting in an overall linear time complexity.

• Space complexity: The space complexity of the solution is $O(n)$, where $n$ is the length of the input array nums. This is because we use additional space to store elements in the “less” and “greater” queues. In the worst case scenario, where all elements are either less or greater than the pivot, the space required would be proportional to the size of the input array nums.