【Go算法速查】快速选择第k个数

package main

import "fmt"

func main() {
    var n, k int
    _, err := fmt.Scanf("%d%d", &n, &k)
    if err != nil {
       return
    }
    arr := make([]int, n)
    for i := 0; i < n; i++ {
       _, err := fmt.Scanf("%d", &arr[i])
       if err != nil {
          return 
       }
    }
    //kthNum := quickSelect(arr, 0, n - 1, k - 1)
    kthNum := quickSelect2(arr, k - 1)
    fmt.Println(kthNum)
}

func quickSelect(arr []int, l int, r int, k int) int {
    if l >= r {
       return arr[k]
    }
    pivot := arr[(l + r) / 2]
    left, right := l, r
    for left <= right {
       for arr[left] < pivot {
          left++
       }
       for arr[right] > pivot {
          right--
       }
       if left <= right {
          arr[left], arr[right] = arr[right], arr[left]
          left++
          right--
       }
    }
    if left <= k {
       return quickSelect(arr, left, r, k)
    }
    if right >= k {
       return quickSelect(arr, l, right, k)
    }
    return arr[k]
}

func quickSelect2(arr []int, k int) int {
    if len(arr) <= 1 {
       return arr[0]
    }
    pivot := arr[len(arr) / 2]
    left, right := 0, len(arr) - 1
    for left <= right {
       for arr[left] < pivot {
          left++
       }
       for arr[right] > pivot {
          right--
       }
       if left <= right {
          arr[left], arr[right] = arr[right], arr[left]
          left++
          right--
       }
    }
    if left <= k {
       // the kth number is in the right part, so we need to adjust the kth index
       return quickSelect2(arr[left:], k - left)
    }
    if right >= k {
       return quickSelect2(arr[: right + 1], k)
    }
    return arr[k]
}