Skip to content

324. Wiggle Sort II

  • Time: $O(n) \to O(n^2)$
  • Space: $O(1)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
 public:
  void wiggleSort(vector<int>& nums) {
    const int n = nums.size();
    const auto it = nums.begin() + n / 2;
    nth_element(nums.begin(), it, nums.end());
    const int median = *it;

// index-rewiring
#define A(i) nums[(1 + 2 * i) % (n | 1)]

    for (int i = 0, j = 0, k = n - 1; i <= k;)
      if (A(i) > median)
        swap(A(i++), A(j++));
      else if (A(i) < median)
        swap(A(i), A(k--));
      else
        ++i;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class Solution {
  public void wiggleSort(int[] nums) {
    final int n = nums.length;
    final int median = findKthLargest(nums, (n + 1) / 2);
    for (int i = 0, j = 0, k = n - 1; i <= k;)
      if (nums[A(i, n)] > median)
        swap(nums, A(i++, n), A(j++, n));
      else if (nums[A(i, n)] < median)
        swap(nums, A(i, n), A(k--, n));
      else
        ++i;
  }

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

  // Same as 215. Kth Largest Element in an Array
  private int findKthLargest(int[] nums, int k) {
    return quickSelect(nums, 0, nums.length - 1, k);
  }

  private int quickSelect(int[] nums, int l, int r, int k) {
    final int pivot = nums[r];

    int nextSwapped = l;
    for (int i = l; i < r; ++i)
      if (nums[i] >= pivot)
        swap(nums, nextSwapped++, i);
    swap(nums, nextSwapped, r);

    final int count = nextSwapped - l + 1; // the number of `nums` >= pivot
    if (count == k)
      return nums[nextSwapped];
    if (count > k)
      return quickSelect(nums, l, nextSwapped - 1, k);
    return quickSelect(nums, nextSwapped + 1, r, k - count);
  }

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class Solution:
  def wiggleSort(self, nums: List[int]) -> None:
    n = len(nums)
    median = self._findKthLargest(nums, (n + 1) // 2)

    def A(i: int):
      return (1 + 2 * i) % (n | 1)

    i = 0
    j = 0
    k = n - 1

    while i <= k:
      if nums[A(i)] > median:
        nums[A(i)], nums[A(j)] = nums[A(j)], nums[A(i)]
        i, j = i + 1, j + 1
      elif nums[A(i)] < median:
        nums[A(i)], nums[A(k)] = nums[A(k)], nums[A(i)]
        k -= 1
      else:
        i += 1

  # Same as 215. Kth Largest Element in an Array
  def _findKthLargest(self, nums: List[int], k: int) -> int:
    def quickSelect(l: int, r: int, k: int) -> int:
      randIndex = random.randint(0, r - l) + l
      nums[randIndex], nums[r] = nums[r], nums[randIndex]
      pivot = nums[r]

      nextSwapped = l
      for i in range(l, r):
        if nums[i] >= pivot:
          nums[nextSwapped], nums[i] = nums[i], nums[nextSwapped]
          nextSwapped += 1
      nums[nextSwapped], nums[r] = nums[r], nums[nextSwapped]

      count = nextSwapped - l + 1  # Number of nums >= pivot
      if count == k:
        return nums[nextSwapped]
      if count > k:
        return quickSelect(l, nextSwapped - 1, k)
      return quickSelect(nextSwapped + 1, r, k - count)

    return quickSelect(0, len(nums) - 1, k)