Update 0215 solution

leetcode/0215.Kth-Largest-Element-in-an-Array/215. Kth Largest Element in an Array.go

@@ -1,6 +1,9 @@
 package leetcode
 
-import "sort"
+import (
+	"math/rand"
+	"sort"
+)
 
 // 解法一 排序，排序的方法反而速度是最快的
 func findKthLargest1(nums []int, k int) int {
@@ -9,36 +12,62 @@ func findKthLargest1(nums []int, k int) int {
 }
 
 // 解法二 这个方法的理论依据是 partition 得到的点的下标就是最终排序之后的下标，根据这个下标，我们可以判断第 K 大的数在哪里
+// 时间复杂度 O(n)，空间复杂度 O(log n)，最坏时间复杂度为 O(n^2)，空间复杂度 O(n)
 func findKthLargest(nums []int, k int) int {
-	if len(nums) == 0 {
-		return 0
-	}
-	return selection(nums, 0, len(nums)-1, len(nums)-k)
+	m := len(nums) - k + 1 // mth smallest, from 1..len(nums)
+	return selectSmallest(nums, 0, len(nums)-1, m)
 }
 
-func selection(arr []int, l, r, k int) int {
-	if l == r {
-		return arr[l]
+func selectSmallest(nums []int, l, r, i int) int {
+	if l >= r {
+		return nums[l]
 	}
-	p := partition164(arr, l, r)
-	if k == p {
-		return arr[p]
-	} else if k < p {
-		return selection(arr, l, p-1, k)
+	q := partition(nums, l, r)
+	k := q - l + 1
+	if k == i {
+		return nums[q]
+	}
+	if i < k {
+		return selectSmallest(nums, l, q-1, i)
 	} else {
-		return selection(arr, p+1, r, k)
+		return selectSmallest(nums, q+1, r, i-k)
 	}
 }
 
-func partition164(a []int, lo, hi int) int {
-	pivot := a[hi]
-	i := lo - 1
-	for j := lo; j < hi; j++ {
-		if a[j] < pivot {
+func partition(nums []int, l, r int) int {
+	k := l + rand.Intn(r-l+1) // 此处为优化，使得时间复杂度期望降为 O(n)，最坏时间复杂度为 O(n^2)
+	nums[k], nums[r] = nums[r], nums[k]
+	i := l - 1
+	// nums[l..i] <= nums[r]
+	// nums[i+1..j-1] > nums[r]
+	for j := l; j < r; j++ {
+		if nums[j] <= nums[r] {
 			i++
-			a[j], a[i] = a[i], a[j]
+			nums[i], nums[j] = nums[j], nums[i]
 		}
 	}
-	a[i+1], a[hi] = a[hi], a[i+1]
+	nums[i+1], nums[r] = nums[r], nums[i+1]
 	return i + 1
 }
+
+// 扩展题 剑指 Offer 40. 最小的 k 个数
+func getLeastNumbers(arr []int, k int) []int {
+	return selectSmallest1(arr, 0, len(arr)-1, k)[:k]
+}
+
+// 和 selectSmallest 实现完全一致，只是返回值不用再截取了，直接返回 nums 即可
+func selectSmallest1(nums []int, l, r, i int) []int {
+	if l >= r {
+		return nums
+	}
+	q := partition(nums, l, r)
+	k := q - l + 1
+	if k == i {
+		return nums
+	}
+	if i < k {
+		return selectSmallest1(nums, l, q-1, i)
+	} else {
+		return selectSmallest1(nums, q+1, r, i-k)
+	}
+}

leetcode/0215.Kth-Largest-Element-in-an-Array/215. Kth Largest Element in an Array_test.go

@@ -27,6 +27,11 @@ func Test_Problem215(t *testing.T) {
 
 	qs := []question215{
 
+		{
+			para215{[]int{3, 2, 1}, 2},
+			ans215{2},
+		},
+
 		{
 			para215{[]int{3, 2, 1, 5, 6, 4}, 2},
 			ans215{5},

website/content/ChapterFour/0200~0299/0215.Kth-Largest-Element-in-an-Array.md

@@ -33,8 +33,8 @@ You may assume k is always valid, 1 ≤ k ≤ array's length.
 
 ## 解题思路
 
-在快排的 partition 操作中，每次 partition 操作结束都会返回一个点，这个标定点的下标和最终排序之后有序数组中这个元素所在的下标是一致的。利用这个特性，我们可以不断的划分数组区间，最终找到第 K 大的元素。执行一次 partition 操作以后，如果这个元素的下标比 K 小，那么接着就在后边的区间继续执行 partition 操作；如果这个元素的下标比 K 大，那么就在左边的区间继续执行 partition 操作；如果相等就直接输出这个下标对应的数组元素即可。
-
+- 在快排的 partition 操作中，每次 partition 操作结束都会返回一个点，这个标定点的下标和最终排序之后有序数组中这个元素所在的下标是一致的。利用这个特性，我们可以不断的划分数组区间，最终找到第 K 大的元素。执行一次 partition 操作以后，如果这个元素的下标比 K 小，那么接着就在后边的区间继续执行 partition 操作；如果这个元素的下标比 K 大，那么就在左边的区间继续执行 partition 操作；如果相等就直接输出这个下标对应的数组元素即可。
+- 快排的思路实现的算法时间复杂度为 O(n)，空间复杂度为 O(logn)。由于证明过程很繁琐，所以不再这里展开讲。具体证明可以参考《算法导论》第 9 章第 2 小节。
 
 
 ## 代码
@@ -43,7 +43,10 @@ You may assume k is always valid, 1 ≤ k ≤ array's length.
 
 package leetcode
 
-import "sort"
+import (
+	"math/rand"
+	"sort"
+)
 
 // 解法一 排序，排序的方法反而速度是最快的
 func findKthLargest1(nums []int, k int) int {
@@ -52,40 +55,66 @@ func findKthLargest1(nums []int, k int) int {
 }
 
 // 解法二 这个方法的理论依据是 partition 得到的点的下标就是最终排序之后的下标，根据这个下标，我们可以判断第 K 大的数在哪里
+// 时间复杂度 O(n)，空间复杂度 O(log n)，最坏时间复杂度为 O(n^2)，空间复杂度 O(n)
 func findKthLargest(nums []int, k int) int {
-	if len(nums) == 0 {
-		return 0
-	}
-	return selection(nums, 0, len(nums)-1, len(nums)-k)
+	m := len(nums) - k + 1 // mth smallest, from 1..len(nums)
+	return selectSmallest(nums, 0, len(nums)-1, m)
 }
 
-func selection(arr []int, l, r, k int) int {
-	if l == r {
-		return arr[l]
+func selectSmallest(nums []int, l, r, i int) int {
+	if l >= r {
+		return nums[l]
 	}
-	p := partition164(arr, l, r)
-	if k == p {
-		return arr[p]
-	} else if k < p {
-		return selection(arr, l, p-1, k)
+	q := partition(nums, l, r)
+	k := q - l + 1
+	if k == i {
+		return nums[q]
+	}
+	if i < k {
+		return selectSmallest(nums, l, q-1, i)
 	} else {
-		return selection(arr, p+1, r, k)
+		return selectSmallest(nums, q+1, r, i-k)
 	}
 }
 
-func partition164(a []int, lo, hi int) int {
-	pivot := a[hi]
-	i := lo - 1
-	for j := lo; j < hi; j++ {
-		if a[j] < pivot {
+func partition(nums []int, l, r int) int {
+	k := l + rand.Intn(r-l+1) // 此处为优化，使得时间复杂度期望降为 O(n)，最坏时间复杂度为 O(n^2)
+	nums[k], nums[r] = nums[r], nums[k]
+	i := l - 1
+	// nums[l..i] <= nums[r]
+	// nums[i+1..j-1] > nums[r]
+	for j := l; j < r; j++ {
+		if nums[j] <= nums[r] {
 			i++
-			a[j], a[i] = a[i], a[j]
+			nums[i], nums[j] = nums[j], nums[i]
 		}
 	}
-	a[i+1], a[hi] = a[hi], a[i+1]
+	nums[i+1], nums[r] = nums[r], nums[i+1]
 	return i + 1
 }
 
+// 扩展题 剑指 Offer 40. 最小的 k 个数
+func getLeastNumbers(arr []int, k int) []int {
+	return selectSmallest1(arr, 0, len(arr)-1, k)[:k]
+}
+
+// 和 selectSmallest 实现完全一致，只是返回值不用再截取了，直接返回 nums 即可
+func selectSmallest1(nums []int, l, r, i int) []int {
+	if l >= r {
+		return nums
+	}
+	q := partition(nums, l, r)
+	k := q - l + 1
+	if k == i {
+		return nums
+	}
+	if i < k {
+		return selectSmallest1(nums, l, q-1, i)
+	} else {
+		return selectSmallest1(nums, q+1, r, i-k)
+	}
+}
+
 ```