LeetCode — 2161. Partition Array According to Given Pivot

This is a medium problem with the description being:

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. If i < j and both elements are smaller (or larger) than pivot, then pi < pj.

Return nums after the rearrangement.

Ok, so for the first and most simple thing that you would imagine, would be to separate in different arrays the small ones, biggest ones and fill up the rest. With that in mind the code would be:

class Solution {
    public int[] pivotArray(int[] nums, int pivot) {

        final List<Integer> beforePivot = new ArrayList<>();
        final List<Integer> afterPivot = new ArrayList<>();
        int amountThatsEqual = 0;

        for(int i=0 ; i < nums.length ; i++) {
            if(nums[i] < pivot) {
                beforePivot.add(nums[i]);
            } else if(nums[i] > pivot) {
                afterPivot.add(nums[i]);
            } else {
                amountThatsEqual++;
            }
        }

        final List<Integer> response = new ArrayList<>();

        response.addAll(beforePivot);

        IntStream.range(0, amountThatsEqual).forEach(i -> response.add(pivot));

        response.addAll(afterPivot);

        return response.stream().mapToInt(Integer::intValue).toArray();
    }
}

Runtime: 25ms, faster than 5.30% of Java online submissions.

Memory Usage: 66.20 MB, less than 63.43% of Java online submissions.

It feels like a good memory usage, but the runtime is really bad. In this first try we use a lot of things that are not needed, and we could go into usage of arrays only. Let’s try to go for just one array and see how it goes:

class Solution {
    public int[] pivotArray(int[] nums, int pivot) {
        final int[] response = new int[nums.length];
        int left = 0;

        for(int i=0 ; i<response.length ; i++) {
            if(nums[i] < pivot) {
                response[left] = nums[i];
                left++;
            }
        }

        for(int i=0 ; i<response.length ; i++) {
        if(nums[i] == pivot) {
                response[left] = pivot;
                left++;
            }
        }

        for(int i=0 ; i<response.length ; i++) {
            if(nums[i] > pivot) {
                response[left] = nums[i];
                left++;
            } 
        }

        return response;
    }
}

Runtime: 5ms, faster than 75.98% of Java online submissions.

Memory Usage: 68.17 MB, less than 33.65% of Java online submissions.

Not bad, runtime much faster, five times faster and we are using only one array, but the bad things is that we are making an iteration three times, and that could be reduced to only one if we used a thing called two pointers technique, which consists into having two markers, usually one at the beginning and the other at the end.

With those you can make your checks and change as needed regarding the question in place.

With that in mind let’s try out:

class Solution {
    public int[] pivotArray(int[] nums, int pivot) {
        final int[] response = new int[nums.length];
        int left = 0;
        int right = nums.length - 1;

        // two pointes one from begining and the other from the end
        for(int i=0,j = nums.length-1 ; i<response.length ; i++, j--) {
            if(nums[i] < pivot) {
                response[left] = nums[i];
                left++;
            }
            if(nums[j] > pivot) {
                response[right] = nums[j];
                right--;
            } 
        }

        // fill the zeros
        while (left <= right) {
            response[left] = pivot;
            left++;
        }

        return response;
    }
}

Runtime: 5ms, faster than 75.98% of Java online submissions.

Memory Usage: 67.28 MB, less than 52.84% of Java online submissions.

Very good, with that we keep good runtime speed, decrease a bit of memory usage and have it more clear and simple.

That’s it! If there is anything thing else to discuss feel free to drop a comment, if I missed anything let me know so I can update accordingly.

Until next post! :)