2161. Partition Array According to Given Pivot

Array
Two Pointers

TIME

O(n)

SPACE

O(n)

Intuition

The result is just three groups concatenated in order: everything < pivot (original order), then all the copies of pivot, then everything > pivot (original order). A stable single pass that buckets each value preserves the required relative order automatically.

Approach

Scan nums once, splitting into two lists while counting the pivot:
- less values strictly less than pivot, appended in order
- greater values strictly greater than pivot, appended in order
- equal_count how many values equal pivot
Build the answer as less + [pivot] * equal_count + greater.

Because we append in scan order, the relative order within less and within greater is preserved.

Complexity

Time: O(n), one pass to bucket, one pass to assemble
Space: O(n), the output array (and the temporary buckets)

Edge Cases

All elements equal the pivot -> result is unchanged
No elements less than (or greater than) the pivot -> that section is empty
Single element (must equal the pivot per constraints) -> returned as-is
Negative values -> handled the same as any other value

Code

from typing import List


class Solution:
    def pivotArray(self, nums: List[int], pivot: int) -> List[int]:
        """
        2161. Partition Array According to Given Pivot
        Time: O(n)
        Space: O(n)
        """
        less = []
        greater = []
        equal_count = 0

        for num in nums:
            if num < pivot:
                less.append(num)
            elif num > pivot:
                greater.append(num)
            else:
                equal_count += 1

        return less + [pivot] * equal_count + greater

/**
 * 2161. Partition Array According to Given Pivot
 * Time: O(n)
 * Space: O(n)
 */
class Solution {
    public int[] pivotArray(int[] nums, int pivot) {
        int n = nums.length;
        int[] result = new int[n];

        int left = 0;
        int right = n - 1;

        // Fill less-than values from the front, greater-than values from the back.
        // Reverse the back fill afterwards to restore relative order.
        for (int i = 0; i < n; i++) {
            if (nums[i] < pivot) {
                result[left++] = nums[i];
            }
            if (nums[n - 1 - i] > pivot) {
                result[right--] = nums[n - 1 - i];
            }
        }

        // Remaining middle slots are the pivot values.
        while (left <= right) {
            result[left++] = pivot;
        }

        return result;
    }
}