Partition Array According to Given Pivot

medium array two pointers simulation

Problem

Given an integer array nums and an integer pivot, rearrange nums so that every element less than pivot comes before every element equal to pivot, which comes before every element greater than pivot. The relative order of equal elements must be preserved.

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

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

Lefts in order: 9, 5, 3. Equals in order: 10, 10. Rights in order: 12, 14.

nums (comma-separated)

pivot

def pivot_array(nums, pivot):
    less, equal, greater = [], [], []
    for x in nums:
        if x < pivot:
            less.append(x)
        elif x == pivot:
            equal.append(x)
        else:
            greater.append(x)
    return less + equal + greater

function pivotArray(nums, pivot) {
  const less = [], equal = [], greater = [];
  for (const x of nums) {
    if (x < pivot) less.push(x);
    else if (x === pivot) equal.push(x);
    else greater.push(x);
  }
  return [...less, ...equal, ...greater];
}

class Solution {
    public int[] pivotArray(int[] nums, int pivot) {
        List<Integer> less = new ArrayList<>();
        List<Integer> equal = new ArrayList<>();
        List<Integer> greater = new ArrayList<>();
        for (int x : nums) {
            if (x < pivot) less.add(x);
            else if (x == pivot) equal.add(x);
            else greater.add(x);
        }
        int[] out = new int[nums.length];
        int i = 0;
        for (int x : less) out[i++] = x;
        for (int x : equal) out[i++] = x;
        for (int x : greater) out[i++] = x;
        return out;
    }
}

vector<int> pivotArray(vector<int>& nums, int pivot) {
    vector<int> less, equal, greater;
    for (int x : nums) {
        if (x < pivot) less.push_back(x);
        else if (x == pivot) equal.push_back(x);
        else greater.push_back(x);
    }
    vector<int> out;
    out.reserve(nums.size());
    for (int x : less) out.push_back(x);
    for (int x : equal) out.push_back(x);
    for (int x : greater) out.push_back(x);
    return out;
}

Explanation

We need every value below pivot first, then the values equal to pivot, then the values above it — and within each group the original order must be preserved. The simplest reliable way to do this is to use three buckets.

We scan the array once and drop each element into the right bucket: less, equal, or greater. Because we visit elements left to right and only ever append, each bucket naturally keeps its items in their original relative order.

At the end we simply concatenate the three lists: less + equal + greater. That guarantees all the small values come first, the pivot copies sit in the middle, and the big values come last.

Example: nums = [9, 12, 5, 10, 14, 3, 10], pivot = 10. The buckets fill as less = [9, 5, 3], equal = [10, 10], greater = [12, 14]. Joined together that gives [9, 5, 3, 10, 10, 12, 14].

It is one clean pass to bucket plus one pass to join, so it runs in linear time using extra arrays for the output.

Time: O(n) Space: O(n)