class Solution(object): def wiggleSort(self, nums): """ :type nums: List[int] :rtype: void Do not return anything, modify nums in-place instead. """ if nums: median = heapq.nsmallest(len(nums) / 2 + 1, nums)[-1] large, small = 1, len(nums) - 1 if len(nums) % 2 == 1 else len(nums) - 2 # go through small section i = small while i > -1: if nums[i] < median: nums[i], nums[small] = nums[small], nums[i] small -= 2 i -= 2 elif nums[i] > median: nums[i], nums[large] = nums[large], nums[i] large += 2 else: i -= 2 # go through large section i = large while i < len(nums): if nums[i] < median: nums[i], nums[small] = nums[small], nums[i] small -= 2 elif nums[i] > median: nums[i], nums[large] = nums[large], nums[i] large += 2 i += 2 else: i += 2
We need to arrange nums into something resembling this:
0 1 2 3 4 5 M S S L L M
We actually need to partition array into 3 sections: large, median and small.
How to do it without virtual indexing?
Let's move the first list towards right:
0 2 4 1 3 5 M S S L L M
Now it is our normal partitioning case. By observation, we know that L starts from 1, so LLM... is 1, 3, 5...; we also know that S starts backwards from len(nums) - 1 if len(nums) is odd, or len(nums) - 2 if len(nums) is even. If we look backwards, SSM is x, x - 2, x - 4...till 0.
I named large and small two pointers to record which part is already built. small pointer means all the numbers after it are smaller than median. large pointer means all the numbers before it are bigger than median.
We can go from the end of S to 0 to build small section first, and then we start from the unbuilt position of large section and build the rest of large section. Then medians must be left in the middle.
# average run time complexity is linear, k is position number def findKthMin(k): left, right = 0, len(nums) - 1 while left < right: pivot = nums[right] pos = left for i in range(left, right): if nums[i] < pivot: nums[pos], nums[i] = nums[i], nums[pos] pos += 1 nums[pos], nums[right] = nums[right], nums[pos] if pos == k: return pivot elif pos > k: right = pos - 1 else: left = pos + 1 return nums[left]
median = findKthMin(len(nums) / 2)
I did write findKthMin using quick select to substitute heapq.nsmalles and OJ accepts it. But it is super slower than the built-in, despite the claim that it is linear running time complexity. I suspect that there must be lots of duplicate numbers. How can I improve quick select with so many duplicates?