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)
}