324. Wiggle Sort II

// LeetCode URL: https://leetcode.com/problems/wiggle-sort-ii/
import java.util.Random;

/**
 * Refer:
 * https://leetcode.com/problems/wiggle-sort-ii/discuss/77682/Step-by-step-explanation-of-index-mapping-in-Java/81849
 *
 * Time Complexity: Best Case O(N). Worst case O(N^2) -- This is due to
 * partition function.
 *
 * Space Complexity: O(N)
 */
class Solution1 {
    public void wiggleSort(int[] nums) throws IllegalArgumentException {
        if (nums == null) {
            throw new IllegalArgumentException("Input nums array is null.");
        }
        int len = nums.length;
        if (len < 2) {
            return;
        }

        int median = findKthLargest(nums, (len + 1) / 2);
        int[] tempArr = new int[len];
        int odd = 1;
        int even = len % 2 == 0 ? len - 2 : len - 1;

        for (int i = 0; i < len; i++) {
            if (nums[i] > median) {
                tempArr[odd] = nums[i];
                odd += 2;
            } else if (nums[i] < median) {
                tempArr[even] = nums[i];
                even -= 2;
            }
        }
        while (odd < nums.length) {
            tempArr[odd] = median;
            odd += 2;
        }
        while (even >= 0) {
            tempArr[even] = median;
            even -= 2;
        }

        for (int i = 0; i < len; i++) {
            nums[i] = tempArr[i];
        }
    }

    private int findKthLargest(int[] nums, int k) {
        shuffle(nums);

        int left = 0;
        int right = nums.length - 1;
        k = nums.length - k;

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

        return nums[left];
    }

    private void shuffle(int[] nums) {
        Random random = new Random();
        for (int i = nums.length - 1; i > 0; i--) {
            int r = random.nextInt(i + 1);
            swap(nums, r, i);
        }
    }

    private int partition(int[] nums, int l, int r) {
        int pivot = l;
        int k = pivot;
        l++;

        while (l <= r) {
            if (nums[l] >= nums[pivot]) {
                k++;
                if (k != l) {
                    swap(nums, k, l);
                }
            }
            l++;
        }

        swap(nums, k, pivot);
        return k;
    }

    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

/*
 * Using 3-way partitioning and virtual indexing. Refer: 1)
 * https://leetcode.com/problems/wiggle-sort-ii/discuss/77677/O(n)+O(1)-after-
 * median-Virtual-Indexing 2)
 * https://leetcode.com/problems/wiggle-sort-ii/discuss/77682/Step-by-step-
 * explanation-of-index-mapping-in-Java
 *
 * Time Complexity: Best Case O(N). Worst case O(N^2) -- This is due to
 * partition function.
 *
 * Space Complexity: O(1)
 */
class Solution2 {
    public void wiggleSort(int[] nums) throws IllegalArgumentException {
        if (nums == null) {
            throw new IllegalArgumentException("Input nums array is null.");
        }
        int len = nums.length;
        if (len < 2) {
            return;
        }

        int median = findKthLargest(nums, (len + 1) / 2);
        int odd = 0;
        int even = len - 1;
        int i = 0;

        while (i <= even) {
            if (nums[nextIdx(i, len)] > median) {
                swap(nums, nextIdx(odd, len), nextIdx(i, len));
                odd++;
                i++;
            } else if (nums[nextIdx(i, len)] < median) {
                swap(nums, nextIdx(even, len), nextIdx(i, len));
                even--;
            } else {
                i++;
            }
        }
    }

    private int nextIdx(int i, int n) {
        return (1 + 2 * i) % (n | 1);
    }

    private int findKthLargest(int[] nums, int k) {
        int left = 0;
        int right = nums.length - 1;
        k = nums.length - k;

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

        return nums[left];
    }

    private int partition(int[] nums, int l, int r) {
        int pivot = l;
        int k = pivot;
        l++;

        while (l <= r) {
            if (nums[l] >= nums[pivot]) {
                k++;
                if (k != l) {
                    swap(nums, k, l);
                }
            }
            l++;
        }

        swap(nums, k, pivot);
        return k;
    }

    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}