Leetcode 2161: Partition array according to given pivot

Question

I'm currently solving LeetCode problem 2161. Partition Array According to Given Pivot:

You are given a 0-indexed integer array nums and an integer pivot. Rearrange nums such that the following conditions are satisfied:

• Every element less than pivot appears before every element greater than pivot.

• Every element equal to pivot appears in between the elements less than and greater than pivot.

• The relative order of the elements less than pivot and the elements greater than pivot is maintained.

  • More formally, consider every pi, pj where pi is the new position of the ith element and pj is the new position of the jth element. For elements less than pivot, if i < j and nums[i] < pivot and nums[j] < pivot, then pi < pj. Similarly for elements greater than pivot, if i < j and nums[i] > pivot and nums[j] > pivot, then pi < pj.

Return nums after the rearrangement.

Example:

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

My solution:

A solution with O(n) memory is trivial. I tried to come up with a O(1) memory solution and approached a similar to insertion sort way:

public int[] pivotArray(int[] nums, int pivot) {
    for(int i = 0; i < nums.length; i++){
        if(nums[i] < pivot){
            for(int j = i; j > 0 && nums[j - 1] >= pivot; j--){
                int tmp = nums[j];
                nums[j] = nums[j - 1];
                nums[j - 1] = tmp;
            }
        }
        if(nums[i] == pivot){
            for(int j = i; j > 0 && nums[j - 1] > pivot; j--){
                int tmp = nums[j];
                nums[j] = nums[j - 1];
                nums[j - 1] = tmp;
            }
        }
    }
    return nums;
}

The issue is the solution has O(n^2) time complexity. Is there an algorithm to solve this faster than O(n^2)? I'm not sure if O(n) time and O(1) space complexity is possible though.

Answer 1

With a slight variation on the answer of Matt Timmermans you could first remove the pivot values from the input array (packing the other values to the left side of it), then apply a O(𝑛log𝑛) divide-and-conquer algorithm to repeatedly merge two subarrays that are already partitioned by the pivot, rotating the center two partitions (of the four partitions). Then finally insert the removed pivot values back between the two partitions that remain.

To avoid consuming O(log𝑛) stack space, perform the divide-and-conquer algorithm in a bottom-up fashion (no recursion).

Here is an implementation:

class Solution {
    void sliceRotateLeft(int[] arr, int start, int end, int shift) {
        // Rotation in O(𝑛), where 𝑛 is the size of the array slice
        int size = end - start;
        if (size <= 1) return;
        shift %= size;
        if (shift == 0) return;
        for (int count = 0, first = 0; count < size; first++, count++) {
            int j = first;
            int value = arr[start + j];
            for (int i = j + shift; i != first; j = i, i = (i + shift) % size) {
                arr[start + j] = arr[start + i];
                count++;
            }
            arr[start + j] = value; 
        }
    }

    int sliceFindPivotIndex(int[] arr, int start, int end, int pivot) {
        // Linear search. NB: binary search would be possible, 
        //                    but does not improve overall complexity.
        while (start < end && arr[start] < pivot) start++;
        return start;
    }

    public int[] pivotArray(int[] nums, int pivot) {
        // Phase 1: pack all non-pivot values to the left
        int len = 0; // Number of non-pivot values
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] != pivot) nums[len++] = nums[i];
        }
        // Phase 2: partition non-pivot values into less-than, greater-than
        // Use bottom up divide-and-conquer
        for (int size = 1; size < len; size *= 2) {
            // Merge two correctly partitioned chunks into one combined chunk
            for (int start = 0; start + size < len; start += size * 2) {
                // Get the offsets of the partitioned subarrays
                int mid = start + size;
                int end = Math.min(mid + size, len);
                int left = sliceFindPivotIndex(nums, start, mid, pivot);
                int right = sliceFindPivotIndex(nums, mid, end, pivot);
                sliceRotateLeft(nums, left, right, mid - left);
            }
        }
        // Phase 3: inject the pivot values
        int pivotFrequency = nums.length - len;
        int i;
        for (i = len - 1; i >= 0 && nums[i] > pivot; i--) {
            nums[i + pivotFrequency] = nums[i];
        }
        for (i++; pivotFrequency > 0; i++, pivotFrequency--) {
            nums[i] = pivot;
        }
        return nums;
    }
}