LeetCode 2161 - Partition Array According to Given Pivot

This problem asks us to rearrange an array around a given pivot value while preserving relative ordering inside certain groups. We are given an integer array nums and an integer pivot. The goal is to reorganize the array into three consecutive sections: 1.

leetcodemediumarraytwo-pointerssimulation

LeetCode Problem 2161

Difficulty: 🟡 Medium
Topics: Array, Two Pointers, Simulation

Solution

Problem Understanding

This problem asks us to rearrange an array around a given pivot value while preserving relative ordering inside certain groups.

We are given an integer array nums and an integer pivot. The goal is to reorganize the array into three consecutive sections:

  1. All elements smaller than pivot
  2. All elements equal to pivot
  3. All elements greater than pivot

The important detail is that the rearrangement must be stable for the elements less than the pivot and for the elements greater than the pivot. Stability means that if two elements originally appeared in a certain order, they must remain in that same order after rearrangement.

For example, if the elements smaller than the pivot appear in the original array as:

9, 5, 3

then they must appear in exactly that order in the final answer. We cannot reorder them into something like:

3, 5, 9

The same rule applies to elements greater than the pivot.

The constraints tell us several important things:

  • The array length can be as large as 10^5
  • Values may be negative
  • The pivot is guaranteed to exist in the array

Since the array can be very large, an inefficient quadratic solution would be too slow. We need a linear or near-linear solution.

There are several edge cases worth thinking about early:

  • Arrays where every element equals the pivot
  • Arrays where all elements are smaller than the pivot
  • Arrays where all elements are greater than the pivot
  • Arrays containing many repeated pivot values
  • Arrays containing negative numbers
  • Arrays of length 1

The problem guarantees that the pivot exists in the array, so we never need to handle the case where the pivot is absent.

Approaches

Brute Force Approach

A brute-force strategy would try to repeatedly move elements into their correct regions by swapping or shifting elements in place.

One possible implementation is:

  • Scan the array from left to right
  • Whenever an element is smaller than the pivot, move it toward the front
  • Whenever an element is greater than the pivot, move it toward the back
  • Carefully shift elements to preserve relative order

This approach can produce the correct answer because it explicitly constructs the desired ordering while maintaining stability.

However, preserving relative order during repeated insertions or shifts is expensive. Every insertion may require moving many elements, leading to a worst-case time complexity of O(n^2).

With n up to 100000, quadratic behavior becomes too slow.

Optimal Approach

The key observation is that the final array is simply the concatenation of three stable groups:

  1. Elements less than the pivot
  2. Elements equal to the pivot
  3. Elements greater than the pivot

We do not actually need complicated in-place rearrangement. Instead, we can perform a linear scan and collect elements into separate lists.

Because we traverse the original array from left to right, the relative ordering inside each group is automatically preserved.

At the end, we concatenate the three lists together to produce the final answer.

This approach is both simple and efficient.

Approach Time Complexity Space Complexity Notes
Brute Force O(n²) O(1) Uses repeated shifting or insertion operations
Optimal O(n) O(n) Stable partition using three temporary lists

Algorithm Walkthrough

Optimal Algorithm

  1. Create three empty lists:
  • less
  • equal
  • greater

These lists will store elements according to their relationship with the pivot. 2. Traverse the array from left to right.

Processing elements in original order is critical because the problem requires stable ordering. 3. For each element:

  • If the value is smaller than the pivot, append it to less
  • If the value equals the pivot, append it to equal
  • If the value is greater than the pivot, append it to greater
  1. After processing all elements:

  • Concatenate the lists in this order:

  • less + equal + greater

  1. Return the resulting array.

Why it works

The algorithm maintains a simple invariant during traversal:

  • less contains all previously seen elements smaller than the pivot, in original order
  • equal contains all previously seen pivot values, in original order
  • greater contains all previously seen elements larger than the pivot, in original order

Since elements are appended in the same order they are encountered, stability is preserved automatically. Concatenating the three groups produces exactly the required final arrangement.

Python Solution

from typing import List

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        less = []
        equal = []
        greater = []

        for value in nums:
            if value < pivot:
                less.append(value)
            elif value == pivot:
                equal.append(value)
            else:
                greater.append(value)

        return less + equal + greater

The implementation directly follows the algorithm described earlier.

We first initialize three lists to represent the three required regions of the final array.

Next, we iterate through nums exactly once. For every value, we compare it with the pivot and append it into the appropriate list. Because appending preserves insertion order, stability is maintained automatically.

Finally, we concatenate the three lists together and return the result.

The code is concise because Python list operations naturally support this type of stable partitioning.

Go Solution

func pivotArray(nums []int, pivot int) []int {
    less := []int{}
    equal := []int{}
    greater := []int{}

    for _, value := range nums {
        if value < pivot {
            less = append(less, value)
        } else if value == pivot {
            equal = append(equal, value)
        } else {
            greater = append(greater, value)
        }
    }

    result := []int{}
    result = append(result, less...)
    result = append(result, equal...)
    result = append(result, greater...)

    return result
}

The Go implementation is conceptually identical to the Python version.

One small difference is that Go does not support direct list concatenation using the + operator for slices. Instead, we use repeated append operations with variadic slice expansion syntax:

append(result, less...)

Go slices dynamically resize as needed, making them a natural fit for this problem.

There are no integer overflow concerns because we only compare and copy values.

Worked Examples

Example 1

Input:

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

We process each element one by one.

Current Value less equal greater
9 [9] [] []
12 [9] [] [12]
5 [9, 5] [] [12]
10 [9, 5] [10] [12]
14 [9, 5] [10] [12, 14]
3 [9, 5, 3] [10] [12, 14]
10 [9, 5, 3] [10, 10] [12, 14]

Final concatenation:

[9,5,3] + [10,10] + [12,14]

Result:

[9,5,3,10,10,12,14]

Example 2

Input:

nums = [-3,4,3,2]
pivot = 2
Current Value less equal greater
-3 [-3] [] []
4 [-3] [] [4]
3 [-3] [] [4, 3]
2 [-3] [2] [4, 3]

Final concatenation:

[-3] + [2] + [4,3]

Result:

[-3,2,4,3]

Complexity Analysis

Measure Complexity Explanation
Time O(n) Each element is processed exactly once
Space O(n) Extra lists store all elements

The algorithm performs a single pass through the array, giving linear time complexity.

The additional space comes from the three temporary lists. Together, they store exactly n elements, so the total auxiliary space complexity is linear.

Test Cases

from typing import List

class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        less = []
        equal = []
        greater = []

        for value in nums:
            if value < pivot:
                less.append(value)
            elif value == pivot:
                equal.append(value)
            else:
                greater.append(value)

        return less + equal + greater

solution = Solution()

# Provided example 1
assert solution.pivotArray(
    [9, 12, 5, 10, 14, 3, 10], 10
) == [9, 5, 3, 10, 10, 12, 14]

# Provided example 2
assert solution.pivotArray(
    [-3, 4, 3, 2], 2
) == [-3, 2, 4, 3]

# Single element array
assert solution.pivotArray(
    [5], 5
) == [5]

# All elements equal to pivot
assert solution.pivotArray(
    [7, 7, 7], 7
) == [7, 7, 7]

# All elements smaller than pivot
assert solution.pivotArray(
    [1, 2, 3], 10
) == [1, 2, 3]

# All elements greater than pivot
assert solution.pivotArray(
    [10, 11, 12], 5
) == [10, 11, 12]

# Multiple pivot occurrences
assert solution.pivotArray(
    [5, 1, 5, 2, 5, 3], 5
) == [1, 2, 3, 5, 5, 5]

# Negative numbers
assert solution.pivotArray(
    [-5, -1, -3, 0, 2], -1
) == [-5, -3, -1, 0, 2]

# Stability check for greater elements
assert solution.pivotArray(
    [8, 6, 7, 5, 3, 0, 9], 5
) == [3, 0, 5, 8, 6, 7, 9]

# Stability check for smaller elements
assert solution.pivotArray(
    [4, 1, 3, 2], 3
) == [1, 2, 3, 4]
Test Why
[9,12,5,10,14,3,10] Validates standard mixed-case behavior
[-3,4,3,2] Validates handling of negative numbers
[5] Tests minimum input size
[7,7,7] Tests all elements equal to pivot
[1,2,3] Tests array with no greater/equal elements
[10,11,12] Tests array with no smaller/equal elements
[5,1,5,2,5,3] Tests multiple pivot occurrences
[-5,-1,-3,0,2] Tests negative values and ordering
[8,6,7,5,3,0,9] Verifies stable ordering for greater elements
[4,1,3,2] Verifies stable ordering for smaller elements

Edge Cases

One important edge case occurs when every element equals the pivot. In this situation, the array should remain unchanged. A buggy implementation might accidentally reorder elements or mishandle empty partitions. Our implementation handles this naturally because all values are appended into the equal list, while the other lists remain empty.

Another important case is when all elements belong entirely to one side of the pivot. For example:

nums = [1,2,3]
pivot = 10

Here, every element is smaller than the pivot. Some implementations may incorrectly assume all three groups contain values. Our solution works correctly because empty lists concatenate safely.

A third important edge case involves preserving stability. Consider:

nums = [8,6,7,5,3,0,9]
pivot = 5

The elements greater than the pivot appear originally as:

8,6,7,9

A non-stable partition algorithm could accidentally reorder them. Our implementation preserves their order because elements are appended sequentially during a left-to-right traversal.

Another subtle case involves repeated pivot values appearing throughout the array. For example:

nums = [5,1,5,2,5,3]

The final arrangement must place all pivot values together while preserving the order of non-pivot groups. The three-list approach guarantees this automatically.