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.

Learn more about Two Pointers patterns.

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.

Solution Implementation

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).

Learn more about how to find time and space complexity quickly using problem constraints.