324. Wiggle Sort II

Wiggle Sort II

Given an unsorted array nums, reorder it such that
nums[0] < nums[1] > nums[2] < nums[3]....



Example:
(1) Given nums = [1, 5, 1, 1, 6, 4], one possible answer is [1, 4, 1, 5, 1, 6]. 
(2) Given nums = [1, 3, 2, 2, 3, 1], one possible answer is [2, 3, 1, 3, 1, 2].



Note:
You may assume all input has valid answer.



Follow Up:
Can you do it in O(n) time and/or in-place with O(1) extra space?

Credits:Special thanks to @dietpepsi for adding this problem and creating all test cases.

Solution

public class Solution {
    public void wiggleSort(int[] nums) {
        if (nums == null || nums.length <= 1) return;

        int size = nums.length;
        int medium = medium(nums);
        int[] newArray = new int[size];
        for (int i = 0; i < size; i++) {
            newArray[i] = medium;
        }

        if (size % 2 == 0) {
            int l = nums.length - 2;
            int r = 1;
            for (int i = 0; i < nums.length; i++) {
                if (nums[i] < medium) {
                    newArray[l] = nums[i];
                    l -= 2;
                } else if (nums[i] > medium) {
                    newArray[r] = nums[i];
                    r += 2;
                }
            }
        } else {
            int i = 0;
            int j = size - 2;
            for (int runner = 0; runner < size; runner++) {
                if (nums[runner] < medium) {
                    newArray[i] = nums[runner];
                    i += 2;
                } else if (nums[runner] > medium) {
                    newArray[j] = nums[runner];
                    j -= 2;
                }
            }
        }

        for (int i = 0; i < size; i++) {
            nums[i] = newArray[i];
        }
    }

    private int medium(int[] nums) {
        int size = nums.length;
        return kTh(nums, size / 2);
    }

    private int kTh(int[] nums, int k) {
        int start = 0;
        int end = nums.length - 1;

        while (start < end) {
            int partition = partition(nums, start, end);
            if (partition == k) {
                return nums[partition];
            } else if (partition < k) {
                start = partition + 1;
            } else {
                end = partition - 1;
            }
        }

        return nums[k];
    }

    private int partition(int[] nums, int start, int end) {
        if (start == end) return start;

        int pivot = nums[start];
        int i = start + 1;
        int j = end;

        while (true) {
            while (i <= end && nums[i] <= pivot) {
                i++;
            }

            while (j >= start + 1 && nums[j] > pivot) {
                j--;
            }

            if (i > j) {
                break;
            }

            swap(nums, i, j);
        }

        swap(nums, start, j);
        return j;
    }

    private void swap(int[] nums, int start, int end) {
        if (start < end) {
            int temp = nums[start];
            nums[start] = nums[end];
            nums[end] = temp;
        }
    }


}