quick sort: top k

Change-Id: I8852538d7fb5df43f39e2493e85bedc68c8da4a2

Author: qieguo2016

Diff for: go/leetcode/215.数组中的第k个最大元素.go

@@ -37,45 +37,86 @@
 
 // @lc code=start
 import (
-	"container/heap"
+	// "container/heap"
+	"math/rand"
 )
 
-type myHeap []int
+// type myHeap []int
 
-func (h *myHeap) Less(i, j int) bool {
-	return (*h)[i] < (*h)[j]
-}
+// func (h *myHeap) Less(i, j int) bool {
+// 	return (*h)[i] < (*h)[j]
+// }
 
-func (h *myHeap) Swap(i, j int) {
-	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
-}
+// func (h *myHeap) Swap(i, j int) {
+// 	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
+// }
+
+// func (h *myHeap) Len() int {
+// 	return len(*h)
+// }
 
-func (h *myHeap) Len() int {
-	return len(*h)
+// func (h *myHeap) Pop() (v interface{}) {
+// 	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
+// 	return
+// }
+
+// func (h *myHeap) Push(v interface{}) {
+// 	*h = append(*h, v.(int))
+// }
+
+func findKthLargest(nums []int, k int) int {
+	// h := make(myHeap, 0)
+	// heap.Init(&h)
+	// for _, cur := range nums {
+	// 	if h.Len() >= k {
+	// 		if cur <= h[0] {
+	// 			continue
+	// 		}
+	// 		heap.Pop(&h)
+	// 	}
+	// 	heap.Push(&h, cur)
+	// }
+	// return h[0]
+	s := &solution{nums: nums}
+	return s.call(len(nums)-k, 0, len(nums)-1)
 }
 
-func (h *myHeap) Pop() (v interface{}) {
-	*h, v = (*h)[:h.Len()-1], (*h)[h.Len()-1]
-	return
+type solution struct {
+	nums []int
 }
 
-func (h *myHeap) Push(v interface{}) {
-	*h = append(*h, v.(int))
+// 用快排的分组思想，如果分组的界限是k，就可以返回了，否则继续2分。
+func (s *solution) call(k, left, right int) int {
+	idx := s.partition(left, right)
+	if k == idx {
+		return s.nums[k]
+	}
+	if k > idx {
+		return s.call(k, idx+1, right)
+	}
+	return s.call(k, left, idx-1)
 }
 
-func findKthLargest(nums []int, k int) int {
-	h := make(myHeap, 0)
-	heap.Init(&h)
-	for _, cur := range nums {
-		if h.Len() >= k {
-			if cur <= h[0] {
-				continue
-			}
-			heap.Pop(&h)
+// 左闭右闭
+func (s *solution) partition(i, j int) int {
+	if i >= j {
+		return i
+	}
+	idx := rand.Intn(j-i) + i
+	s.nums[idx], s.nums[i] = s.nums[i], s.nums[idx]
+	d := s.nums[i]
+	for i < j {
+		for i < j && s.nums[j] >= d {
+			j--
+		}
+		s.nums[i] = s.nums[j]
+		for i < j && s.nums[i] < d {
+			i++
 		}
-		heap.Push(&h, cur)
+		s.nums[j] = s.nums[i]
 	}
-	return h[0]
+	s.nums[i] = d
+	return i
 }
 
 // @lc code=end