Partition Array According to Given Pivot

Initially, I went for a 3 pointer approach; just count the number of elements that is smaller than the pivot, l, and elements that are equal to the pivot, s. The 3 pointers we are initializing are left = 0, middle = l and right = l + s.

Then, after initializing an ans array, we can re-iterate through nums and:

• add all elements that are smaller than the pivot to ans[left] and increment left by 1
• add all elements that are equal to the pivot to ans[middle] and increment middle by 1
• add all elements that are larget to the pivot to ans[right] and increment right by 1

code

class Solution {
public:
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        int l = 0;
        int s = 0;
        for (int n : nums) {
            if (n < pivot) {
                ++l;
            } else if (n == pivot) {
                ++s;
            }
        }
        int l_index = 0;
        int s_index = l;
        int m_index = l+s;
        vector<int> new_nums(nums.size());
        for (int n : nums) {
            if (n < pivot) {
                new_nums[l_index] = n;
                l_index++;
            } else if (n == pivot) {
                new_nums[s_index] = n;
                s_index++;
            } else {
                new_nums[m_index] = n;
                m_index++;
            }
        }
        return new_nums;
    }
};

However, the above solution requires 2 traversals of nums.

In fact, we can shorten it down by simultaneously iterating from both sides at the same time. Initializing an ans array, and keeping track of 2 variables start and end, we know that:

• when iterating from left to right, we add elements smaller than the pivot to ans[start] and increment start
• when iterating from right to left, we add elements larger than the pivot to ans[end] and decrement end

Then, the remaining elements between index start and end must be equal to pivot.

code

class Solution {
public:
    vector<int> pivotArray(vector<int>& nums, int pivot) {
        int beginning = 0;
        int ending = nums.size()-1;

        vector<int> ans(nums.size());
        for (int i = 0, j = nums.size()-1; i < nums.size(), j >= 0; ++i, --j) {
            if (nums[i] < pivot) {
                ans[beginning] = nums[i];
                ++beginning;
            }
            if (nums[j] > pivot) {
                ans[ending] = nums[j];
                --ending;
            }
        }
        while(beginning <= ending) {
            ans[beginning++] = pivot;
        }
        return ans;
    }
};

complexity

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