u

top-100-liked/73/leetcode_test.go

@@ -0,0 +1,48 @@
+package top100liked
+
+import (
+	"math/rand"
+	"testing"
+)
+
+// https://leetcode.cn/problems/kth-largest-element-in-an-array/?envType=study-plan-v2&envId=top-100-liked
+
+func findKthLargest(nums []int, k int) int {
+	n := len(nums)
+	target := n - k // 第 k 大 = 升序第 n-k 位
+	var quickSelect func(l, r int) int
+	quickSelect = func(l, r int) int {
+		if l == r {
+			return nums[l]
+		}
+		pivot := nums[l+rand.Intn(r-l+1)]
+		// 三路分区：[l..lt-1] < pivot, [lt..gt] == pivot, [gt+1..r] > pivot
+		lt, gt, i := l, r, l
+		for i <= gt {
+			switch {
+			case nums[i] < pivot:
+				nums[lt], nums[i] = nums[i], nums[lt]
+				lt++
+				i++
+			case nums[i] > pivot:
+				nums[gt], nums[i] = nums[i], nums[gt]
+				gt--
+			default:
+				i++
+			}
+		}
+		switch {
+		case target < lt:
+			return quickSelect(l, lt-1)
+		case target > gt:
+			return quickSelect(gt+1, r)
+		default:
+			return pivot
+		}
+	}
+	return quickSelect(0, n-1)
+}
+
+func Test(t *testing.T) {
+	findKthLargest([]int{3, 2, 1, 5, 6, 4}, 2)
+}