2161. Partition Array According to Given Pivot


Problem Description

In this problem, we are given an array nums and an integer pivot. The goal is to rearrange the array such that:

  1. All elements less than pivot are positioned before all elements greater than pivot.
  2. All elements equal to pivot are placed in the middle, between the elements less than and those greater than pivot.
  3. The relative order of elements less than and those greater than pivot must be the same as in the original array.

In summary, the task is to partition the array into three parts: the first containing elements less than pivot, the second containing elements equal to pivot, and the third containing elements greater than pivot, while preserving the original relative order of the elements in the first and third parts.

Intuition

To solve this problem, we can approach it by separating the elements into three distinct lists based on their comparison to the pivot:

  1. A list to hold elements less than pivot.
  2. A list for elements equal to pivot.
  3. A list for elements greater than pivot.

After segregating the elements into the respective lists, we can then concatenate these lists in the order of less than pivot, equal to pivot, and greater than pivot. This concatenation will result in the desired array that fulfills all the problem conditions.

The reason we use separate lists instead of in-place swaps is that in-place operations might make it complex to preserve the original relative order. Simple list operations like appending and concatenation keep the original order intact and make the implementation straightforward and efficient.

This approach ensures that we only pass through the array once, making the algorithm linear in time because each element is considered exactly once and placed into one of the three lists.

Solution Approach

The solution is implemented in Python and uses a simple and effective algorithm involving basic list operations. Here's the walk-through of the implementation:

  1. Three separate lists are initialized: a to hold elements less than pivot, b for elements equal to pivot, and c for elements greater than pivot.

    1a, b, c = [], [], []
  2. The algorithm proceeds by iterating through each element x in the given array nums.

    1for x in nums:
  3. For each element x, a comparison is made to classify it into one of the three lists:

    • If x is less than pivot, it is appended to the list a.
    • If x is equal to pivot, it is appended to list b.
    • If x is greater than pivot, it is appended to list c.
    1if x < pivot:
    2    a.append(x)
    3elif x == pivot:
    4    b.append(x)
    5else:
    6    c.append(x)
  4. By the end of the loop, all the elements are distributed among the three lists, preserving their original relative order within each category (less than, equal to, and greater than pivot).

  5. The final step is to concatenate the three lists: a (elements less than pivot), b (elements equal to pivot), and c (elements greater than pivot). This results in the rearranged array that meets all the required conditions.

    1return a + b + c

Through the use of lists and the built-in list method append, the solution takes advantage of Python's dynamic array capabilities. This eliminates the need for complex index management or in-place replacements that might compromise the relative order of the elements.

The solution relies on the efficiency of Python's underlying memory management for dynamic arrays, and it works within the confines of O(n) space complexity (where n is the number of elements in nums) because it creates separate lists for partitioning the data, which are later merged. The time complexity is also O(n), as each element is looked at exactly once during the for-loop iteration.

Example Walkthrough

Let's consider a small example to illustrate the solution approach. Suppose we have the following array and pivot:

1nums = [9, 12, 3, 5, 14, 10, 10]
2pivot = 10

Now, we will apply the solution algorithm step by step:

  1. Initialize three empty lists a, b, and c to categorize the elements as less than, equal to, and greater than the pivot, respectively:

    1a = []
    2b = []
    3c = []
  2. Iterate through each element x in the array nums:

    • We start with x = 9 which is less than the pivot, so we append it to the list a: a = [9].
    • Next, x = 12 is greater than the pivot, appended to c: c = [12].
    • Then, x = 3 is less than the pivot, appended to a: a = [9, 3].
    • Followed by x = 5 which is again less than the pivot, so a becomes a = [9, 3, 5].
    • We proceed to x = 14 which is greater than the pivot and append it to c: c = [12, 14].
    • Next, we have two elements equal to the pivot, x = 10, so we append both to b: b = [10, 10].

    At the end of the iteration, the lists are as follows:

    1a = [9, 3, 5]
    2b = [10, 10]
    3c = [12, 14]
  3. All the elements have been classified into three separate lists while preserving their original order within each list category.

  4. We concatenate the three lists in the order of a, b, and c to get the final result:

    1result = a + b + c
    2# result = [9, 3, 5, 10, 10, 12, 14]

    The final array is [9, 3, 5, 10, 10, 12, 14] which satisfies the condition of keeping all elements less than 10 before all elements equal to 10 and those greater than 10, while preserving the original relative order within each category.

By using this approach, we've maintained a simple, understandable, and efficient solution that neatly classifies and recombines the elements as desired.

Python Solution

1class Solution:
2    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
3        # Initialize lists to hold numbers smaller than,
4        # equal to, and greater than the pivot
5        smaller_than_pivot = []
6        equal_to_pivot = []
7        greater_than_pivot = []
8
9        # Iterate through each number in the input list
10        for number in nums:
11            # If the number is less than the pivot,
12            # add it to the smaller_than_pivot list
13            if number < pivot:
14                smaller_than_pivot.append(number)
15            # If the number is equal to the pivot,
16            # add it to the equal_to_pivot list
17            elif number == pivot:
18                equal_to_pivot.append(number)
19            # If the number is greater than the pivot,
20            # add it to the greater_than_pivot list
21            else:
22                greater_than_pivot.append(number)
23
24        # Combine the lists and return the result,
25        # with all numbers less than the pivot first,
26        # followed by numbers equal to the pivot,
27        # and finally numbers greater than the pivot
28        return smaller_than_pivot + equal_to_pivot + greater_than_pivot
29

Java Solution

1class Solution {
2    // This method takes an array 'nums' and an integer 'pivot', then reorders the array such that
3    // all elements less than 'pivot' come before elements equal to 'pivot', and those come before elements greater than 'pivot'.
4    public int[] pivotArray(int[] nums, int pivot) {
5        int n = nums.length; // Get the length of the array.
6        int[] ans = new int[n]; // Create a new array 'ans' to store the reordered elements.
7        int index = 0; // Initialize an index variable to keep track of the position in 'ans' array.
8      
9        // First pass: Place all elements less than 'pivot' into the 'ans' array.
10        for (int num : nums) {
11            if (num < pivot) {
12                ans[index++] = num;
13            }
14        }
15      
16        // Second pass: Place all elements equal to 'pivot' into the 'ans' array.
17        for (int num : nums) {
18            if (num == pivot) {
19                ans[index++] = num;
20            }
21        }
22      
23        // Third pass: Place all elements greater than 'pivot' into the 'ans' array.
24        for (int num : nums) {
25            if (num > pivot) {
26                ans[index++] = num;
27            }
28        }
29      
30        return ans; // Return the reordered array.
31    }
32}
33

C++ Solution

1#include <vector>
2
3class Solution {
4public:
5    // Function to rearrange elements in an array with respect to a pivot element.
6    // All elements less than pivot come first, followed by elements equal to pivot,
7    // and then elements greater than pivot.
8    std::vector<int> pivotArray(std::vector<int>& nums, int pivot) {
9        // Vector to store the rearranged elements.
10        std::vector<int> rearranged;
11
12        // First pass: add elements less than pivot to rearranged vector.
13        for (int num : nums) {
14            if (num < pivot) {
15                rearranged.push_back(num);
16            }
17        }
18      
19        // Second pass: add elements equal to pivot to rearranged vector.
20        for (int num : nums) {
21            if (num == pivot) {
22                rearranged.push_back(num);
23            }
24        }
25
26        // Third pass: add elements greater than pivot to rearranged vector.
27        for (int num : nums) {
28            if (num > pivot) {
29                rearranged.push_back(num);
30            }
31        }
32
33        // Return the vector containing elements in the desired order.
34        return rearranged;
35    }
36};
37

Typescript Solution

1// Import the array class from TypeScript default library
2import { number } from "prop-types";
3
4// Function to rearrange elements in an array with respect to a pivot element.
5// All elements less than pivot come first, followed by elements equal to pivot,
6// and then elements greater than pivot.
7function pivotArray(nums: number[], pivot: number): number[] {
8    // Array to store the rearranged elements.
9    let rearranged: number[] = [];
10
11    // First pass: add elements less than pivot to rearranged array.
12    for (let num of nums) {
13        if (num < pivot) {
14            rearranged.push(num);
15        }
16    }
17  
18    // Second pass: add elements equal to pivot to rearranged array.
19    for (let num of nums) {
20        if (num === pivot) {
21            rearranged.push(num);
22        }
23    }
24
25    // Third pass: add elements greater than pivot to rearranged array.
26    for (let num of nums) {
27        if (num > pivot) {
28            rearranged.push(num);
29        }
30    }
31
32    // Return the array containing elements in the desired order.
33    return rearranged;
34}
35

Time and Space Complexity

Time Complexity:

The time complexity of the given code relies primarily on the for loop that iterates over all n elements in the input list nums. Inside the loop, the code compares each element with the pivot and then adds it to one of the three lists (a, b, or c). Each of these operations—comparison and append—is performed in constant time, O(1). However, combining the lists at the end a + b + c takes O(n) time since it creates a new list containing all n elements. Therefore, the overall time complexity of the code is O(n).

Space Complexity:

The space complexity refers to the amount of extra space or temporary space used by the algorithm. In this case, we're creating three separate lists (a, b, and c) to hold elements less than, equal to, and greater than the pivot, respectively. In the worst-case scenario, all elements could be less than, equal to, or greater than the pivot, leading to each list potentially containing n elements. Therefore, the additional space used by the lists is directly proportional to the number of elements in the input, n. Thus, the space complexity is O(n).


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫