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?

Solution

public class Solution {

    int n;
    public void wiggleSort(int[] nums) {
        n = nums.length;
        int median = getMedian(nums);
        // three way partition based on index mapping
        int i = 0, j = 0, k = n - 1;
        while(j <= k) {
            // we examine each j
            if(nums[mapping(j)] < median) {
                swap(nums, mapping(j), mapping(i));
                j++;
                i++;
            } else if(nums[mapping(j)] > median) {
                swap(nums, mapping(j), mapping(k));
                k--;
            } else {
                j++;
            }
        }

    }

    public int mapping(int i) {
        if(i < (n+1)/2) return n % 2 == 0 ? n - 2 * i - 2 : n - 2 * i - 1;
        else return (2 * n - 2 * i - 1);
    }

    public void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }

    public int getMedian(int[] nums) {
        int k = nums.length / 2;

        int l = 0;
        int r = nums.length - 1;
        int idx = partition(nums, l, r);
        while(idx != k) {
            if(idx > k) {
                r = idx - 1;
            } else {
                l = idx + 1;
            }
            idx = partition(nums, l, r);
        }
        return nums[idx];
    }

    public int partition(int[] nums, int l, int r) {
        if(l == r) return l;

        int pivot = nums[r];
        int pos = l; // every element before pos is < pivot
        for(int i = l; i <= r; i++) {
            if(nums[i] < pivot) {
                swap(nums, pos++, i);
            }
        }
        swap(nums, pos, r);
        return pos;
    }
}