Golang quick select solution


  • 1

    Let's say we have an array nums as below and k = 2:
    [4, 7, 5, 9, 2, 3]

    First we chose a pivot. For simplicity, I choose nums[0] = 4. You can also choose nums[end] or nums[mid], if you want.
    After choosing the pivot, we organize the array so that the array looks like [(less than 4), 4, (greater than 4)].
    If we perform a normal quick sort, then we will recursively run the same code for the former half and the latter half of nums.

    But for solving this problem, we should stop here and think what is happening now.
    We can't guarantee the left side and the right side is sorted, but we can guarantee all numbers left side of 4 is less than 4 and all ones right side of 4 is greater than 4.
    Then, if there is only 1 number at the right side of 4, can't we say 4 is k = 2 nd largest number in nums?

    And even if 4 is not k th largest number in nums, like binary search, we need to only either side of 4 for next quick sort. We aren't interested in sorting, we just are interested in which number is k th largest.
    So, before finding k th element, we need to only swap k times elements. That is to say, we can achieve expected O(k) complexity, while in worst case it's still O(n^2).

    func findKthLargest(nums []int, k int) int {
        return doFindKthLargest(nums, k, 0, len(nums)-1)
    }
    
    func doFindKthLargest(nums []int, k int, start, end int) int {
        nlen := len(nums)
        targetPos := nlen - k
        
        for {
            pivotIndex := partition(nums, start, end)
            if pivotIndex == targetPos {
                return nums[pivotIndex]
            } else if pivotIndex < targetPos {
                start = pivotIndex + 1
            } else {
                end = pivotIndex - 1
            }
        }
    }
    
    // partition choses nums[start] as pivot and
    // moves all values less than pivot to the left side,
    // moves all values greater than pivot to the rigit side
    // of the pivot value.
    // This represents an "one shot" implementation of quick sort.
    func partition(nums []int, start, end int) int {
        pivot := nums[start]
        left, right := start+1, end
        
        for left <= right {
            for left <= right && nums[left] <= pivot {
                left++
            }
            for left <= right && nums[right] > pivot {
                right--
            }
            if left <= right {
                nums[left], nums[right] = nums[right], nums[left]
            }
        }
        nums[right], nums[start] = nums[start], nums[right]
        return right
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.