Clear Java O(n) avg time & O(n) space solution using 3-way-partition


  • 16

    The basic idea is to first select the kth smallest element, where k is the half the size (if size is even) or half the size+1 (if size is odd). So the array is partitioned into two halves, for even size, left half size==right half size, for odd size, left half size==right half size+1. Then iterate back from left and right half, put each value into original array.

    For example, [1, 5, 1, 1, 6, 4], after select kth smallest element, it becomes [1,1,1,5,6,4] with median index of 2. For the left half is [1,1,1], right half is [5,6,4]. After merge, it becomes, 1,4,1,6,1,5.

    Same for [4,5,5,6], after select kth smallest , it becomes [4,5,5,6] with left half [4,5] and right half [5,6], merge it to be [5,6,4,5].

    The tricky parts:

    1. Do a three-way-partition, so that the duplicates are around median, so that they can be split

    2. Iterate from the back of two halves, so that the duplicates in between can be split apart.

      public class Solution {
      public void wiggleSort(int[] nums) {
      int median = selectKth(nums, 0, nums.length-1, nums.length%2==0 ? nums.length/2 : nums.length/2+1);
      List<Integer> leftArr = new ArrayList<Integer>();
      for(int i=0; i<=median; i++)
      leftArr.add(nums[i]);
      List<Integer> rightArr = new ArrayList<Integer>();
      for(int i=median+1; i<nums.length; i++)
      rightArr.add(nums[i]);
      for(int li=leftArr.size()-1,ri=rightArr.size()-1,i=0; ri>=0; li--,ri--,i+=2) { // right is same or shorter than left
      nums[i] = leftArr.get(li);
      nums[i+1] = rightArr.get(ri);
      }
      if(nums.length%2!=0)
      nums[nums.length-1] = leftArr.get(0);
      }

       private int selectKth(int[] nums, int start, int end, int k) {
           int[] res = partition(nums,start,end);
           int lb = res[0]; int hb = res[1];
           if(k-1<lb)
               return selectKth(nums,start,lb-1,k);
           else if (k-1>hb)
               return selectKth(nums,hb+1,end,k);
           else
               return k-1;
       }
       
       private int[] partition(int[] nums, int lb, int hb) {
           int pVal = nums[lb]; // use random genarater is better in performance
           int i = lb;
           while(i<=hb) {
               if(nums[i]==pVal)
                   i++;
               else if(nums[i]<pVal)
                   swap(nums,i++,lb++);
               else
                   swap(nums,i,hb--);
           }
           int[] res = new int[2];
           res[0] = lb; res[1] = hb;
           return res;
       }
       
       private void swap(int[] nums, int i, int j) {
           int tmp = nums[i];
           nums[i] = nums[j];
           nums[j] = tmp;
       }
      

      }


  • 0
    T

    I guess for finding kth smalelst element might take O(nlog n) complexity?


  • 0
    C

    finding kth smalelst element is O(n) for expectation


  • 0
    C

    How do you handle such a case: leftArr=[2,2,1,1] rightArr=[2,2,3,3], the answer seems to be [2,2,2,2,1,3,1,3], which is not right. I'm stuck here.


  • 0

    Hi, i think your leftArr is not correct, after selecting kth smallest, the leftArr = [1,1,2,2], so after merge, it is [2,3,2,3,1,2,1,2], which is correct


  • 0

    And one more thing, when merge the two arrays, we need to iterate from BACK of leftArr and rightArr, this is actually a tricky part to split over duplicates.


  • 0
    C

    Well i've noticed that there is a trick in partition, swapping ones smaller than pVal to the beginning, and only ones larger than pVal to the end, making pVal never present in the beginning or the end if possible. But, I could not figure out why it works to iterate from back, could you please give more hints?


  • 1

    This is based on observations, even though I cannot prove it.
    After kth selection, we know that if there are duplicates involved between leftArr and rightArr, they are within the last few elements of leftArr and first few elements of rightArr.
    For example, leftArr = [1....5,5] while rightArr is [5,5,7,.....]. The last few elements of rightArr normally is larger than the last few elements of leftArr. Iterate from back of leftArr and rightArr will help proceed the duplicates of leftArr (in this case, 5) as soon as possible. Just take [4,5,5,6] as an example, you will know, what happens if you iterate from start and what is the difference if you iterate from back.


  • 0
    C

    Assume that there are n//2 elements equaling median and x of them in arrLeft, so the last element equaling median from arrLeft will be placed at 2x. It can be worked out that there are x elements larger than median, taking places from 1 to 2x+1, so elements equaling median from arrRight will be placed from 2x+3, which will never be beside 2x. Another situation, n is odd and there are there are n//2+1 elements equaling median; if no element smaller than median, this method works; if there are, no solution.


  • 0
    X

    your tricky parts are really nice !


  • 0
    N

    @pinkfloyda Your solution fails on following test case
    Input: [1,1,2,1,2,2,1]
    Expected: [1,2,1,2,1,2,1]
    Output: [1,2,1,2,1,2,2]


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.